"""SB3 + Decision Transformer baselines for the ReMDM diffusion planner. This module wraps standard discrete-action RL baselines (PPO, A2C, DQN, recurrent PPO) plus two imitation baselines (Behavioural Cloning and Decision Transformer) into the project's unified config + dispatch surface so they can be compared head-to-head against the DAgger / offline-BC diffusion planner on the same MiniHack environments. Entry point: :func:`run_baselines`. Hyperparameters live in ``configs/defaults.yaml`` under the ``baselines_*`` namespace; the unified env-step training budget (``cfg.total_timesteps``) is shared with DAgger and offline BC. W&B logging routes through the project's :class:`Logger` (with the W&B project temporarily swapped to ``cfg.baselines_wandb_project``); SB3's standard ``WandbCallback`` piggybacks on the active run and syncs its tensorboard scalars automatically. No file in this module calls ``wandb.log(...)`` directly. """ from __future__ import annotations import importlib.util import logging import os import random import zlib from pathlib import Path from types import SimpleNamespace from typing import Any import gymnasium as gym import numpy as np import orjson import torch import torch.nn as nn from sb3_contrib import RecurrentPPO from stable_baselines3 import A2C, DQN, PPO from stable_baselines3.common.callbacks import CallbackList, EvalCallback from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.policies import ActorCriticPolicy from stable_baselines3.common.torch_layers import BaseFeaturesExtractor from stable_baselines3.common.vec_env import SubprocVecEnv from torch.utils.data import DataLoader, Dataset from wandb.integration.sb3 import WandbCallback from src.envs.minihack_env import ( AdvancedObservationEnv, collect_oracle_trajectory, ) from src.planners.logging import Logger logger = logging.getLogger(__name__) SB3_RL_ALGOS: tuple[str, ...] = ("ppo", "a2c", "dqn", "ppo-rnn") IMITATION_ALGOS: tuple[str, ...] = ("bc", "dt") ALL_BASELINE_ALGOS: tuple[str, ...] = SB3_RL_ALGOS + IMITATION_ALGOS class _SB3MiniHackWrapper(gym.Wrapper): """Reshape ``AdvancedObservationEnv`` tuple obs into an SB3 dict obs. The underlying env returns ``(local_crop, global_map)`` with shapes ``(crop, crop)`` and ``(map_h, map_w)``; SB3's ``MultiInputPolicy`` needs a ``Dict`` space with explicit channel dims. Also remaps ``info["won"]`` -> ``info["is_success"]`` so SB3's success tracking reports our win rate. """ def __init__(self, env: AdvancedObservationEnv) -> None: super().__init__(env) local_h, local_w = env.observation_space.shape cfg = env._cfg # AdvancedObservationEnv stores cfg here self.observation_space = gym.spaces.Dict( { "local": gym.spaces.Box( low=0, high=6000, shape=(1, local_h, local_w), dtype=np.int16, ), "global": gym.spaces.Box( low=0, high=6000, shape=(1, cfg.map_h, cfg.map_w), dtype=np.int16, ), } ) def reset(self, **kwargs: Any) -> tuple[dict[str, np.ndarray], dict]: (local, glob), info = self.env.reset(**kwargs) return self._pack(local, glob), info def step( self, action: int, ) -> tuple[dict[str, np.ndarray], float, bool, bool, dict]: (local, glob), reward, terminated, truncated, info = self.env.step(action) if "won" in info: info["is_success"] = info["won"] return self._pack(local, glob), reward, terminated, truncated, info @staticmethod def _pack( local: np.ndarray, glob: np.ndarray, ) -> dict[str, np.ndarray]: return { "local": np.expand_dims(local, axis=0), # [1, crop, crop] "global": np.expand_dims(glob, axis=0), # [1, H, W] } class _MiniHackCNN(BaseFeaturesExtractor): """Dual-stream CNN for the SB3 dict observation. Local stream: ``Conv(1->16, 3) -> Conv(16->32, 3)``. Global stream: ``Conv(1->16, 5, stride 2) -> Conv(16->32, 3, stride 2)``. Both streams are flattened and concatenated, then projected to ``features_dim`` via a single linear + ReLU. """ def __init__( self, observation_space: gym.spaces.Dict, features_dim: int = 256, ) -> None: super().__init__(observation_space, features_dim) self.local_cnn = nn.Sequential( nn.Conv2d(1, 16, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(), nn.Flatten(), ) self.global_cnn = nn.Sequential( nn.Conv2d(1, 16, kernel_size=5, stride=2), nn.ReLU(), nn.Conv2d(16, 32, kernel_size=3, stride=2), nn.ReLU(), nn.Flatten(), ) with torch.no_grad(): dummy_loc = torch.zeros(1, *observation_space["local"].shape) dummy_glob = torch.zeros(1, *observation_space["global"].shape) n_flatten = ( self.local_cnn(dummy_loc).shape[1] + self.global_cnn(dummy_glob).shape[1] ) self.linear = nn.Sequential(nn.Linear(n_flatten, features_dim), nn.ReLU()) def forward( self, observations: dict[str, torch.Tensor], ) -> torch.Tensor: loc = self.local_cnn(observations["local"].float()) # [B, F_l] glob = self.global_cnn(observations["global"].float()) # [B, F_g] return self.linear(torch.cat([loc, glob], dim=1)) class _MiniHackStateEncoder(nn.Module): """CNN encoder mapping a (local, global) obs pair to a state embedding.""" def __init__( self, embed_dim: int = 128, crop_h: int = 9, crop_w: int = 9, map_h: int = 21, map_w: int = 79, ) -> None: super().__init__() self.local_cnn = nn.Sequential( nn.Conv2d(1, 16, kernel_size=3, padding=1), nn.ReLU(), nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(), nn.Flatten(), ) self.global_cnn = nn.Sequential( nn.Conv2d(1, 16, kernel_size=5, stride=2), nn.ReLU(), nn.Conv2d(16, 32, kernel_size=3, stride=2), nn.ReLU(), nn.Flatten(), ) with torch.no_grad(): dummy_loc = torch.zeros(1, 1, crop_h, crop_w) dummy_glob = torch.zeros(1, 1, map_h, map_w) local_flat = self.local_cnn(dummy_loc).shape[1] global_flat = self.global_cnn(dummy_glob).shape[1] self.proj = nn.Linear(local_flat + global_flat, embed_dim) def forward( self, local_obs: torch.Tensor, global_obs: torch.Tensor, ) -> torch.Tensor: # Accepts (B, T, 1, H, W) or (B, 1, H, W). if local_obs.dim() == 5: B, T = local_obs.shape[:2] local_obs = local_obs.view(B * T, *local_obs.shape[2:]) global_obs = global_obs.view(B * T, *global_obs.shape[2:]) reshape = True else: B, T = local_obs.shape[0], 1 reshape = False loc_feat = self.local_cnn(local_obs.float()) # [B*T, F_l] glob_feat = self.global_cnn(global_obs.float()) # [B*T, F_g] out = self.proj(torch.cat([loc_feat, glob_feat], dim=-1)) # [B*T, D] if reshape: out = out.view(B, T, -1) return out class _DecisionTransformer(nn.Module): """Causal Decision Transformer over interleaved (R, s, a) tokens.""" def __init__( self, n_actions: int, embed_dim: int = 128, n_heads: int = 4, n_layers: int = 3, context_len: int = 30, max_ep_len: int = 500, dropout: float = 0.1, crop_h: int = 9, crop_w: int = 9, map_h: int = 21, map_w: int = 79, ) -> None: super().__init__() self.embed_dim = embed_dim self.context_len = context_len self.n_actions = n_actions self.max_ep_len = max_ep_len self.state_encoder = _MiniHackStateEncoder( embed_dim, crop_h, crop_w, map_h, map_w, ) self.action_embed = nn.Embedding(n_actions + 1, embed_dim) # +1 for pad self.return_embed = nn.Linear(1, embed_dim) self.pos_embed = nn.Embedding(max_ep_len, embed_dim) self.token_type_embed = nn.Embedding(3, embed_dim) self.embed_ln = nn.LayerNorm(embed_dim) self.dropout = nn.Dropout(dropout) encoder_layer = nn.TransformerEncoderLayer( d_model=embed_dim, nhead=n_heads, dim_feedforward=embed_dim * 4, dropout=dropout, activation="gelu", batch_first=True, ) self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers) self.action_head = nn.Linear(embed_dim, n_actions) self.apply(self._init_weights) @staticmethod def _init_weights(module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): nn.init.normal_(module.weight, mean=0.0, std=0.02) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) def forward( self, returns_to_go: torch.Tensor, # [B, T, 1] local_obs: torch.Tensor, # [B, T, 1, H_l, W_l] global_obs: torch.Tensor, # [B, T, 1, H_g, W_g] actions: torch.Tensor, # [B, T] timesteps: torch.Tensor, # [B, T] attention_mask: torch.Tensor | None = None, # [B, T] ) -> torch.Tensor: B, T = returns_to_go.shape[:2] device = returns_to_go.device rtg_embed = self.return_embed(returns_to_go) # [B, T, D] state_embed = self.state_encoder(local_obs, global_obs) # [B, T, D] action_embed = self.action_embed(actions) # [B, T, D] pos_embed = self.pos_embed(timesteps) # [B, T, D] rtg_embed = rtg_embed + pos_embed + self.token_type_embed.weight[0] state_embed = state_embed + pos_embed + self.token_type_embed.weight[1] action_embed = action_embed + pos_embed + self.token_type_embed.weight[2] # Interleave (R_0, s_0, a_0, R_1, s_1, a_1, ...) -> [B, 3T, D] stacked = torch.stack([rtg_embed, state_embed, action_embed], dim=2) stacked = stacked.view(B, 3 * T, self.embed_dim) stacked = self.dropout(self.embed_ln(stacked)) seq_len = 3 * T causal_mask = torch.triu( torch.ones(seq_len, seq_len, device=device), diagonal=1, ).bool() key_padding_mask = None if attention_mask is not None: expanded = attention_mask.unsqueeze(-1).repeat(1, 1, 3).view(B, 3 * T) key_padding_mask = expanded == 0 hidden = self.transformer( stacked, mask=causal_mask, src_key_padding_mask=key_padding_mask, ) # State token positions are 1, 4, 7, ... -> stride 3. state_hidden = hidden[:, 1::3, :] # [B, T, D] return self.action_head(state_hidden) # [B, T, A] @torch.no_grad() def get_action( self, returns_to_go: torch.Tensor, local_obs: torch.Tensor, global_obs: torch.Tensor, actions: torch.Tensor, timesteps: torch.Tensor, ) -> torch.Tensor: self.eval() logits = self.forward( returns_to_go, local_obs, global_obs, actions, timesteps, ) return logits[:, -1, :].argmax(dim=-1) class _DTDataset(Dataset): """Sliding-window dataset over Decision Transformer trajectories.""" def __init__( self, trajectories: list[dict[str, np.ndarray]], context_len: int, max_ep_len: int, n_actions: int, ) -> None: self.trajectories = trajectories self.context_len = context_len self.max_ep_len = max_ep_len self.n_actions = n_actions self.indices: list[tuple[int, int]] = [ (traj_idx, start) for traj_idx, traj in enumerate(trajectories) for start in range(len(traj["actions"])) ] def __len__(self) -> int: return len(self.indices) def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: traj_idx, start = self.indices[idx] traj = self.trajectories[traj_idx] traj_len = len(traj["actions"]) end = min(start + self.context_len, traj_len) actual_len = end - start local = traj["local"][start:end].copy() glob = traj["global"][start:end].copy() actions = traj["actions"][start:end].copy() rtg = traj["returns_to_go"][start:end].copy() timesteps = np.arange(start, end) # Clamp to valid embedding ranges. timesteps = np.clip(timesteps, 0, self.max_ep_len - 1) actions = np.clip(actions, 0, self.n_actions - 1) pad_len = self.context_len - actual_len if pad_len > 0: local = np.pad( local, ((0, pad_len), (0, 0), (0, 0), (0, 0)), mode="constant", ) glob = np.pad( glob, ((0, pad_len), (0, 0), (0, 0), (0, 0)), mode="constant", ) actions = np.pad(actions, (0, pad_len), mode="constant") rtg = np.pad(rtg, (0, pad_len), mode="constant") timesteps = np.pad(timesteps, (0, pad_len), mode="constant") attention_mask = np.zeros(self.context_len, dtype=np.float32) attention_mask[:actual_len] = 1.0 return { "local": torch.tensor(local, dtype=torch.float32), "global": torch.tensor(glob, dtype=torch.float32), "actions": torch.tensor(actions, dtype=torch.long), "returns_to_go": torch.tensor(rtg, dtype=torch.float32).unsqueeze(-1), "timesteps": torch.tensor(timesteps, dtype=torch.long), "attention_mask": torch.tensor(attention_mask, dtype=torch.float32), } class _PrefixedEvalCallback(EvalCallback): """``EvalCallback`` that records mean_reward / avg_steps / win_rate under a unique per-environment prefix. SB3 truncates metric names at 36 chars, which collides on long MiniHack env IDs; the prefix lets us strip ``MiniHack-`` / ``-v0`` cleanly. """ def __init__( self, eval_env: SubprocVecEnv, prefix: str, **kwargs: Any, ) -> None: super().__init__(eval_env, **kwargs) self.prefix = prefix def _on_step(self) -> bool: cont = super()._on_step() if self.evaluations_results: self.logger.record( f"{self.prefix}/mean_reward", float(np.mean(self.evaluations_results[-1])), ) self.logger.record( f"{self.prefix}/avg_steps", float(np.mean(self.evaluations_length[-1])), ) if self.evaluations_successes: self.logger.record( f"{self.prefix}/win_rate", float(np.mean(self.evaluations_successes[-1])), ) return cont def quiet_multiprocessing_tempdir_teardown() -> None: """Stop multiprocessing printing a traceback it cannot act on at exit. ``SubprocVecEnv`` makes multiprocessing create a ``pymp-*`` directory under ``TMPDIR`` and register a finalizer that ``shutil.rmtree``s it at interpreter exit. That finalizer tolerates ``FileNotFoundError`` and re-raises everything else, and on a shared filesystem an unlinked-but-still-open file survives as ``.nfsXXXX`` until its last handle closes, so the rmtree raises ``OSError: Directory not empty``. ``multiprocessing.util._run_finalizers`` catches that and prints the traceback, which is why ``--mode baselines`` ends with a traceback on stderr, exit code 0 and every artefact written. This widens the finalizer's tolerance to any removal error. Nothing is left behind that was not already: the directory is inside the process's own temp root, and the leftover ``.nfsXXXX`` entries are reaped by the filesystem once the last handle closes. Must run before the first ``SubprocVecEnv``: multiprocessing captures the callback when it creates the directory, so a later patch is ignored. Idempotent, and a no-op if a future Python drops the private hook. """ import multiprocessing.util as mp_util original = getattr(mp_util, "_remove_temp_dir", None) if original is None or getattr(original, "_remdm_tolerant", False): return def tolerant_remove_temp_dir(rmtree, tempdir): # `**kwargs` swallows the `onerror`/`onexc` handler the caller # supplies, whose whole job is to re-raise; `ignore_errors` replaces # it and takes precedence in `shutil.rmtree` regardless. def quiet_rmtree(path, **kwargs): rmtree(path, ignore_errors=True) try: original(quiet_rmtree, tempdir) except OSError as exc: # pragma: no cover - filesystem dependent logger.debug("Left multiprocessing temp dir %s: %s", tempdir, exc) tolerant_remove_temp_dir._remdm_tolerant = True mp_util._remove_temp_dir = tolerant_remove_temp_dir def _make_sb3_env_fn(env_id: str, cfg: SimpleNamespace, log_dir: str): """Return a picklable thunk that builds one wrapped+monitored env.""" def _init() -> Monitor: os.makedirs(log_dir, exist_ok=True) env = AdvancedObservationEnv(env_id, des_file=None, cfg=cfg) env = _SB3MiniHackWrapper(env) return Monitor(env, log_dir) return _init def _short(env_id: str) -> str: return env_id.replace("MiniHack-", "").replace("-v0", "") def _eval_episodes_per_env(cfg: SimpleNamespace) -> int: override = getattr(cfg, "baselines_eval_episodes_per_env", None) if override is not None: return int(override) return int(cfg.eval_episodes_per_env) def _seed_everything(seed: int) -> None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) def _resolve_output_dir(cfg: SimpleNamespace, override: str | None) -> Path: out = Path(override) if override else Path(cfg.baselines_output_dir) out.mkdir(parents=True, exist_ok=True) return out def _init_baseline_logger( cfg: SimpleNamespace, run_name: str, ) -> Logger: """Init the project Logger with W&B project swapped to baselines. Mutates ``cfg.wandb_project`` / ``cfg.wandb_run_name`` / ``cfg.wandb_resume_id`` for the duration of the call so the existing Logger constructor picks them up. We deliberately do not restore the originals — each baseline seed reuses this helper, and main.py exits after ``run_baselines`` returns. """ project_override = getattr(cfg, "baselines_wandb_project", None) if project_override: cfg.wandb_project = project_override cfg.wandb_run_name = run_name cfg.wandb_resume_id = None return Logger(cfg) def _collect_bc_dataset( cfg: SimpleNamespace, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Roll out the BFS oracle on each ID env and stack flat (s, a) pairs.""" n_per_env = int(cfg.baselines_bc_oracle_episodes_per_env) locals_, globals_, actions_ = [], [], [] for env_id in cfg.id_envs: for traj_seed in range(n_per_env): traj = collect_oracle_trajectory(env_id, traj_seed, cfg) if traj is None: continue # (T, H, W) -> (T, 1, H, W) locals_.append(np.expand_dims(traj["local"], axis=1)) globals_.append(np.expand_dims(traj["global"], axis=1)) actions_.append(traj["actions"]) if not actions_: raise RuntimeError("BC oracle collection produced zero trajectories") return ( np.concatenate(locals_, axis=0), np.concatenate(globals_, axis=0), np.concatenate(actions_, axis=0), ) class _BCDataset(Dataset): def __init__( self, loc: np.ndarray, glob: np.ndarray, acts: np.ndarray, ) -> None: self.loc = torch.tensor(loc, dtype=torch.float32) self.glob = torch.tensor(glob, dtype=torch.float32) self.acts = torch.tensor(acts, dtype=torch.int64) def __len__(self) -> int: return len(self.acts) def __getitem__( self, idx: int, ) -> dict[str, dict[str, torch.Tensor] | torch.Tensor]: return { "obs": {"local": self.loc[idx], "global": self.glob[idx]}, "acts": self.acts[idx], } def evaluation_seeds(env_id: str, n_episodes: int) -> list[int]: """The per-episode evaluation seeds, matching the planner's exactly. The planner evaluates on `42 + crc32(f"{env_id}:{ep}") % 2**31` (`src/planners/inference.py`, pinned by `test_evaluator_seeds_are_fixed_and_run_seed_independent`). The baselines have to evaluate on the same episodes or the headline planner-vs-baseline comparison is between two different sets of levels. They did not. `_make_sb3_env_fn` built the env with no seed at all and `_eval_sb3_policy_manually` ran it inside a `SubprocVecEnv`, a child process that never inherited the parent's `_seed_everything`; `_eval_dt` had the same shape. The only seeding was the Python/NumPy/torch globals plus `seed=` into the SB3 constructors, which seeds action sampling and **not MiniHack level generation** — gymnasium's `reset(seed=...)` does not reach the NetHack core RNG, which is why `AdvancedObservationEnv.reset` seeds it explicitly. Measured before the fix: `_seed_everything(0)` twice gave first-observation hashes `62a012aa6f073246` and `2ee59c01b9bfd35d`. Python's `hash()` is salted per process, so crc32 is what keeps these stable across invocations. The formula is duplicated from the planner rather than imported, and `test_baseline_eval_seeds_match_the_planners` fails if the two ever drift. Args: env_id: The MiniHack environment id being evaluated. n_episodes: How many episodes the evaluation runs. Returns: One seed per episode, in episode order. """ return [ 42 + zlib.crc32(f"{env_id}:{ep}".encode()) % (2**31) for ep in range(n_episodes) ] def _eval_sb3_policy_manually( policy: ActorCriticPolicy, env_id: str, cfg: SimpleNamespace, log_dir: str, n_episodes: int, ) -> tuple[float, float]: """Run ``policy.predict`` on a Monitor-wrapped env and return (win_rate, avg_steps). Each episode is reset on its own seed from :func:`evaluation_seeds`, so the baseline is scored on the same levels the planner is. That is also why the env is built here rather than inside a ``SubprocVecEnv``: the vec env auto- resets between episodes with no seed to give it, and it ran in a child process that never inherited the parent's ``_seed_everything`` — one env in one subprocess bought nothing and cost the seeding. """ os.makedirs(log_dir, exist_ok=True) eval_env = Monitor( _SB3MiniHackWrapper(AdvancedObservationEnv(env_id, des_file=None, cfg=cfg)), log_dir, ) seeds = evaluation_seeds(env_id, n_episodes) try: wins = 0 total_steps = 0 for episode in range(n_episodes): obs, _info = eval_env.reset(seed=seeds[episode]) terminated = truncated = False info: dict = {} while not (terminated or truncated): action, _ = policy.predict(obs, deterministic=True) obs, _reward, terminated, truncated, info = eval_env.step(int(action)) if info.get("won", False): wins += 1 total_steps += info["episode"]["l"] finally: eval_env.close() return wins / n_episodes, total_steps / n_episodes def _train_bc( cfg: SimpleNamespace, train_env: SubprocVecEnv, log: Logger, log_dir: str, seed: int, ) -> tuple[ActorCriticPolicy, dict[str, float]]: """Train a Behavioural Cloning baseline. Returns (policy, seed_metrics).""" device = torch.device(cfg.device) n_eval = _eval_episodes_per_env(cfg) logger.info("Collecting oracle demonstrations for BC...") loc_arr, glob_arr, acts_arr = _collect_bc_dataset(cfg) logger.info("BC dataset: %d transitions", len(acts_arr)) bc_loader = DataLoader( _BCDataset(loc_arr, glob_arr, acts_arr), batch_size=int(cfg.baselines_bc_batch_size), shuffle=True, num_workers=4, pin_memory=torch.cuda.is_available(), ) lr = float(cfg.baselines_bc_lr) policy = ActorCriticPolicy( observation_space=train_env.observation_space, action_space=train_env.action_space, lr_schedule=lambda _progress: lr, features_extractor_class=_MiniHackCNN, features_extractor_kwargs={"features_dim": 256}, ).to(device) n_epochs = int(cfg.baselines_bc_epochs) optimizer = torch.optim.AdamW( policy.parameters(), lr=lr, weight_decay=float(cfg.weight_decay), ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=n_epochs, ) policy.train() for epoch in range(n_epochs): total_loss = 0.0 for batch in bc_loader: obs = {k: v.to(policy.device) for k, v in batch["obs"].items()} acts = batch["acts"].to(policy.device) _values, log_prob, _entropy = policy.evaluate_actions(obs, acts) loss = -log_prob.mean() optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(policy.parameters(), 1.0) optimizer.step() total_loss += loss.item() scheduler.step() avg_loss = total_loss / max(1, len(bc_loader)) current_lr = scheduler.get_last_lr()[0] log.log( { "train/bc_loss": avg_loss, "train/lr": current_lr, "train/epoch": epoch + 1, }, step=epoch + 1, ) logger.info( "BC epoch %02d/%02d | loss=%.4f | lr=%.2e", epoch + 1, n_epochs, avg_loss, current_lr, ) seed_metrics: dict[str, float] = {} for split, env_list in (("ID", cfg.id_envs), ("OOD", cfg.ood_envs)): logger.info("--- BC %s evaluation (seed=%d) ---", split, seed) for env_id in env_list: short = _short(env_id) win_rate, avg_steps = _eval_sb3_policy_manually( policy, env_id, cfg, f"{log_dir}/eval_{split.lower()}/{env_id}", n_eval, ) seed_metrics[f"{split}/{short}/win_rate"] = win_rate * 100 seed_metrics[f"{split}/{short}/avg_steps"] = avg_steps logger.info( "%-30s | win_rate=%5.1f%% | avg_steps=%5.1f", short, win_rate * 100, avg_steps, ) log.log(seed_metrics, step=n_epochs + 1) return policy, seed_metrics def _collect_dt_trajectories( cfg: SimpleNamespace, ) -> list[dict[str, np.ndarray]]: """Collect oracle trajectories with sparse reward + return-to-go labels.""" n_per_env = int(cfg.baselines_dt_oracle_episodes_per_env) trajectories: list[dict[str, np.ndarray]] = [] for env_id in cfg.id_envs: for traj_seed in range(n_per_env): traj = collect_oracle_trajectory(env_id, traj_seed, cfg) if traj is None: continue T = len(traj["actions"]) rewards = np.zeros(T, dtype=np.float32) rewards[-1] = 1.0 # sparse goal reward rtg = np.zeros(T, dtype=np.float32) rtg[-1] = rewards[-1] for t in range(T - 2, -1, -1): rtg[t] = rewards[t] + rtg[t + 1] trajectories.append( { "local": np.expand_dims(traj["local"], axis=1), "global": np.expand_dims(traj["global"], axis=1), "actions": traj["actions"], "rewards": rewards, "returns_to_go": rtg, } ) return trajectories def _eval_dt( model: _DecisionTransformer, env_id: str, cfg: SimpleNamespace, target_return: float, n_episodes: int, max_ep_len: int, eval_max_steps: int, context_len: int, ) -> tuple[float, float]: """Roll out a trained Decision Transformer with target-return conditioning.""" device = torch.device(cfg.device) env = AdvancedObservationEnv(env_id, des_file=None, cfg=cfg) env = _SB3MiniHackWrapper(env) # The same episodes the planner and the SB3 baselines are scored on. seeds = evaluation_seeds(env_id, n_episodes) model.eval() wins = 0 total_steps = 0 try: for _ep in range(n_episodes): obs, _ = env.reset(seed=seeds[_ep]) done = False local_hist: list[np.ndarray] = [] global_hist: list[np.ndarray] = [] action_hist: list[int] = [] rtg_hist: list[float] = [] ts_hist: list[int] = [] current_rtg = float(target_return) t = 0 info: dict = {} while not done and t < eval_max_steps: local_hist.append(obs["local"]) global_hist.append(obs["global"]) rtg_hist.append(current_rtg) ts_hist.append(min(t, max_ep_len - 1)) ctx = min(len(local_hist), context_len) local_in = np.stack(local_hist[-ctx:], axis=0) global_in = np.stack(global_hist[-ctx:], axis=0) rtg_in = np.array(rtg_hist[-ctx:], dtype=np.float32) ts_in = np.array(ts_hist[-ctx:], dtype=np.int64) if len(action_hist) < ctx: act_in = np.zeros(ctx, dtype=np.int64) if action_hist: act_in[-len(action_hist) :] = action_hist[-ctx:] else: act_in = np.array(action_hist[-ctx:], dtype=np.int64) local_t = ( torch.tensor(local_in, dtype=torch.float32).unsqueeze(0).to(device) ) global_t = ( torch.tensor(global_in, dtype=torch.float32).unsqueeze(0).to(device) ) rtg_t = ( torch.tensor(rtg_in, dtype=torch.float32) .unsqueeze(0) .unsqueeze(-1) .to(device) ) act_t = torch.tensor(act_in, dtype=torch.long).unsqueeze(0).to(device) ts_t = torch.tensor(ts_in, dtype=torch.long).unsqueeze(0).to(device) with torch.no_grad(): action = int( model.get_action(rtg_t, local_t, global_t, act_t, ts_t).item() ) action = max(0, min(action, int(cfg.action_dim) - 1)) action_hist.append(action) obs, reward, terminated, truncated, info = env.step(action) done = terminated or truncated current_rtg -= float(reward) t += 1 if info.get("won", False): wins += 1 total_steps += t finally: env.close() return wins / n_episodes, total_steps / n_episodes def _train_dt( cfg: SimpleNamespace, log: Logger, log_dir: str, seed: int, ) -> tuple[_DecisionTransformer, dict[str, float]]: """Train a Decision Transformer baseline. Returns (model, seed_metrics).""" device = torch.device(cfg.device) context_len = int(cfg.baselines_dt_context_len) max_ep_len = int(cfg.baselines_dt_max_ep_len) eval_max_steps = int(cfg.baselines_dt_eval_max_steps) n_eval = _eval_episodes_per_env(cfg) n_epochs = int(cfg.baselines_dt_epochs) logger.info("Collecting oracle demonstrations for DT...") trajectories = _collect_dt_trajectories(cfg) if not trajectories: raise RuntimeError("DT oracle collection produced zero trajectories") traj_lengths = [len(t["actions"]) for t in trajectories] logger.info( "DT dataset: %d trajectories, %d transitions (len: min=%d max=%d mean=%.1f)", len(trajectories), sum(traj_lengths), min(traj_lengths), max(traj_lengths), float(np.mean(traj_lengths)), ) if max(traj_lengths) > max_ep_len: logger.warning( "Longest oracle trajectory (%d) exceeds baselines_dt_max_ep_len (%d); " "positions will be clamped.", max(traj_lengths), max_ep_len, ) target_return = float(max(t["returns_to_go"][0] for t in trajectories)) dataset = _DTDataset( trajectories, context_len=context_len, max_ep_len=max_ep_len, n_actions=int(cfg.action_dim), ) loader = DataLoader( dataset, batch_size=int(cfg.baselines_dt_batch_size), shuffle=True, num_workers=4, pin_memory=torch.cuda.is_available(), ) model = _DecisionTransformer( n_actions=int(cfg.action_dim), embed_dim=int(cfg.baselines_dt_embed_dim), n_heads=int(cfg.baselines_dt_n_heads), n_layers=int(cfg.baselines_dt_n_layers), context_len=context_len, max_ep_len=max_ep_len, crop_h=int(cfg.crop_size), crop_w=int(cfg.crop_size), map_h=int(cfg.map_h), map_w=int(cfg.map_w), ).to(device) n_params = sum(p.numel() for p in model.parameters()) logger.info("DT parameters: %d", n_params) optimizer = torch.optim.AdamW( model.parameters(), lr=float(cfg.baselines_dt_lr), weight_decay=float(cfg.weight_decay), ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=n_epochs, ) for epoch in range(n_epochs): model.train() total_loss = 0.0 n_batches = 0 for batch in loader: local = batch["local"].to(device) glob = batch["global"].to(device) actions = batch["actions"].to(device) rtg = batch["returns_to_go"].to(device) timesteps = batch["timesteps"].to(device) attention_mask = batch["attention_mask"].to(device) logits = model(rtg, local, glob, actions, timesteps, attention_mask) logits_flat = logits.reshape(-1, int(cfg.action_dim)) targets_flat = actions.reshape(-1) mask_flat = attention_mask.reshape(-1) ce = nn.functional.cross_entropy( logits_flat, targets_flat, reduction="none", ) loss = (ce * mask_flat).sum() / mask_flat.sum().clamp(min=1.0) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss += loss.item() n_batches += 1 scheduler.step() avg_loss = total_loss / max(1, n_batches) log.log( { "train/dt_loss": avg_loss, "train/lr": float(scheduler.get_last_lr()[0]), "train/epoch": epoch + 1, }, step=epoch + 1, ) logger.info( "DT epoch %02d/%02d | loss=%.4f | lr=%.2e", epoch + 1, n_epochs, avg_loss, float(scheduler.get_last_lr()[0]), ) seed_metrics: dict[str, float] = {} logger.info("DT eval target return = %.2f", target_return) for split, env_list in (("ID", cfg.id_envs), ("OOD", cfg.ood_envs)): logger.info("--- DT %s evaluation (seed=%d) ---", split, seed) for env_id in env_list: short = _short(env_id) win_rate, avg_steps = _eval_dt( model, env_id, cfg, target_return=target_return, n_episodes=n_eval, max_ep_len=max_ep_len, eval_max_steps=eval_max_steps, context_len=context_len, ) seed_metrics[f"{split}/{short}/win_rate"] = win_rate * 100 seed_metrics[f"{split}/{short}/avg_steps"] = avg_steps logger.info( "%-30s | win_rate=%5.1f%% | avg_steps=%5.1f", short, win_rate * 100, avg_steps, ) log.log(seed_metrics, step=n_epochs + 1) return model, seed_metrics def _build_sb3_model( algo: str, train_env: SubprocVecEnv, cfg: SimpleNamespace, seed: int, tb_log_dir: str, ): """Construct one of {ppo, a2c, dqn, ppo-rnn} with the MiniHack CNN.""" policy_kwargs = { "features_extractor_class": _MiniHackCNN, "features_extractor_kwargs": {"features_dim": 256}, } # SB3 raises at learn() time if tensorboard_log is set without tensorboard # installed, and tensorboard is not a dependency of this project. if not importlib.util.find_spec("tensorboard"): logger.info("tensorboard not installed; SB3 tensorboard logging off") tb_log_dir = None if algo == "ppo": return PPO( "MultiInputPolicy", train_env, policy_kwargs=policy_kwargs, verbose=1, tensorboard_log=tb_log_dir, seed=seed, ) if algo == "ppo-rnn": return RecurrentPPO( "MultiInputLstmPolicy", train_env, policy_kwargs=policy_kwargs, verbose=1, tensorboard_log=tb_log_dir, seed=seed, ) if algo == "a2c": return A2C( "MultiInputPolicy", train_env, policy_kwargs=policy_kwargs, verbose=1, tensorboard_log=tb_log_dir, seed=seed, ) if algo == "dqn": return DQN( "MultiInputPolicy", train_env, policy_kwargs=policy_kwargs, verbose=1, tensorboard_log=tb_log_dir, seed=seed, buffer_size=int(cfg.baselines_dqn_buffer_size), ) raise ValueError(f"Unknown SB3 algo: {algo!r}") def _build_sb3_callbacks( cfg: SimpleNamespace, train_env: SubprocVecEnv, log_dir: str, model_dir: str, ) -> CallbackList: # WandbCallback requires an active run; without one it raises on # construction, which would make every SB3 baseline unrunnable with # use_wandb=false. import wandb callbacks: list = [] if wandb.run is not None: callbacks.append(WandbCallback(model_save_path=model_dir)) n_eval = _eval_episodes_per_env(cfg) eval_freq = max( 1, int(cfg.baselines_eval_freq_env_steps) // train_env.num_envs, ) for env_id in cfg.id_envs: short = _short(env_id) eval_env = SubprocVecEnv( [_make_sb3_env_fn(env_id, cfg, f"{log_dir}/eval_id/{env_id}")] ) callbacks.append( _PrefixedEvalCallback( eval_env, prefix=f"ID/{short}", best_model_save_path=f"{model_dir}/best_{env_id}/", log_path=f"{log_dir}/eval_id/{env_id}/", eval_freq=eval_freq, n_eval_episodes=n_eval, deterministic=True, ) ) for env_id in cfg.ood_envs: short = _short(env_id) eval_env = SubprocVecEnv( [_make_sb3_env_fn(env_id, cfg, f"{log_dir}/eval_ood/{env_id}")] ) callbacks.append( _PrefixedEvalCallback( eval_env, prefix=f"OOD/{short}", best_model_save_path=None, log_path=f"{log_dir}/eval_ood/{env_id}/", eval_freq=eval_freq, n_eval_episodes=n_eval, deterministic=True, ) ) return CallbackList(callbacks) def _aggregate( all_seed_results: list[dict[str, Any]], ) -> dict[str, dict[str, float | list[float]]]: """Compute mean/std across seeds for every shared metric key.""" if not all_seed_results: return {} metric_keys = [k for k in all_seed_results[0] if k != "seed"] agg: dict[str, dict[str, float | list[float]]] = {} for key in metric_keys: values = [r[key] for r in all_seed_results if key in r] if values: agg[key] = { "mean": float(np.mean(values)), "std": float(np.std(values)), "values": [float(v) for v in values], } return agg def _print_aggregated(seeds: list[int], agg: dict[str, dict[str, Any]]) -> None: if not agg: logger.info( "No per-environment metrics to aggregate (RL eval is callback-driven)" ) return logger.info("Aggregated results across %d seeds: %s", len(seeds), seeds) for split in ("ID", "OOD"): env_metrics: dict[str, dict[str, dict[str, Any]]] = {} for key, stats in agg.items(): if not key.startswith(f"{split}/"): continue _split, env_name, metric_name = key.split("/", 2) env_metrics.setdefault(env_name, {})[metric_name] = stats if not env_metrics: continue logger.info("--- %s environments ---", split) for env_name, metrics in sorted(env_metrics.items()): wr = metrics.get("win_rate", {}) steps = metrics.get("avg_steps", {}) logger.info( "%-30s | win_rate=%5.1f%% +/- %4.1f | avg_steps=%5.1f +/- %4.1f", env_name, wr.get("mean", 0.0), wr.get("std", 0.0), steps.get("mean", 0.0), steps.get("std", 0.0), ) def _save_aggregated( out_path: Path, algo: str, seeds: list[int], all_seed_results: list[dict[str, Any]], agg: dict[str, dict[str, Any]], ) -> None: payload = { "algorithm": algo, "seeds": seeds, "n_seeds": len(seeds), "per_seed_results": all_seed_results, "aggregated": {k: {"mean": v["mean"], "std": v["std"]} for k, v in agg.items()}, } out_path.write_bytes(orjson.dumps(payload, option=orjson.OPT_INDENT_2)) logger.info("Aggregated results written to %s", out_path) def run_baselines( cfg: SimpleNamespace, algo: str, seeds: list[int] | None = None, output_path: str | None = None, ) -> None: """Train and evaluate one baseline algorithm across one or more seeds. Args: cfg: Project config namespace (must contain ``baselines_*`` keys). algo: One of ``ppo``, ``a2c``, ``dqn``, ``ppo-rnn``, ``bc``, ``dt``. seeds: Optional list of seeds. ``None`` -> ``[cfg.seed]`` (or a single seed of ``0`` if ``cfg.seed`` is ``None``). output_path: Optional override for the aggregated-results JSON destination. When ``None``, results land under ``cfg.baselines_output_dir``. """ if algo not in ALL_BASELINE_ALGOS: raise ValueError(f"Unknown algo {algo!r}. Choose one of {ALL_BASELINE_ALGOS}.") # Before the first SubprocVecEnv, or the finalizer is already registered. quiet_multiprocessing_tempdir_teardown() if seeds is None: seeds = [cfg.seed if cfg.seed is not None else 0] if not seeds: raise ValueError("seeds must be non-empty") out_dir = _resolve_output_dir(cfg, None) if output_path is not None: agg_json_path = Path(output_path) agg_json_path.parent.mkdir(parents=True, exist_ok=True) else: agg_json_path = out_dir / f"results_{algo}_{len(seeds)}seeds.json" logger.info( "Running baseline %s on %d seed(s): %s (output -> %s)", algo, len(seeds), seeds, agg_json_path, ) all_seed_results: list[dict[str, Any]] = [] n_envs_per_id = int(cfg.baselines_n_envs_per_id) for seed_idx, seed in enumerate(seeds): logger.info( "============================================================\n" " %s seed %d (%d/%d)\n" "============================================================", algo.upper(), seed, seed_idx + 1, len(seeds), ) _seed_everything(seed) run_name = f"{algo}-multitask-seed{seed}" log = _init_baseline_logger(cfg, run_name) run_id = ( log._run.id # type: ignore[union-attr] if log._use_wandb and log._run is not None else f"local-{algo}-seed{seed}" ) log_dir = str(out_dir / "logs" / run_id) model_dir = str(out_dir / "models" / run_id) os.makedirs(log_dir, exist_ok=True) os.makedirs(model_dir, exist_ok=True) seed_results: dict[str, Any] = {"seed": seed} try: if algo == "dt": model, dt_metrics = _train_dt(cfg, log, log_dir, seed) seed_results.update(dt_metrics) torch.save( { "model_state_dict": model.state_dict(), "config": { "n_actions": int(cfg.action_dim), "embed_dim": int(cfg.baselines_dt_embed_dim), "n_heads": int(cfg.baselines_dt_n_heads), "n_layers": int(cfg.baselines_dt_n_layers), "context_len": int(cfg.baselines_dt_context_len), "max_ep_len": int(cfg.baselines_dt_max_ep_len), }, }, f"{model_dir}/dt_final_seed{seed}.pt", ) else: # SB3 RL families and BC both need the parallel train env. train_env_fns = [ _make_sb3_env_fn(env_id, cfg, log_dir) for env_id in list(cfg.id_envs) * n_envs_per_id ] train_env = SubprocVecEnv(train_env_fns) try: if algo == "bc": policy, bc_metrics = _train_bc( cfg, train_env, log, log_dir, seed, ) seed_results.update(bc_metrics) policy.save(f"{model_dir}/bc_final_seed{seed}") else: sb3_model = _build_sb3_model( algo, train_env, cfg, seed, tb_log_dir=str(out_dir / "tb" / run_id), ) callbacks = _build_sb3_callbacks( cfg, train_env, log_dir, model_dir, ) logger.info( "Training %s for %d env-steps across %d ID maps " "(%d parallel envs)...", algo.upper(), int(cfg.total_timesteps), len(cfg.id_envs), train_env.num_envs, ) sb3_model.learn( total_timesteps=int(cfg.total_timesteps), callback=callbacks, ) sb3_model.save(f"{model_dir}/{algo}_final_seed{seed}") finally: train_env.close() all_seed_results.append(seed_results) finally: log.finish() logger.info("%s seed %d complete.", algo.upper(), seed) agg = _aggregate(all_seed_results) _print_aggregated(seeds, agg) if agg: _save_aggregated(agg_json_path, algo, seeds, all_seed_results, agg) # Final summary write to the project Logger so the aggregated # numbers land on a dedicated W&B run. summary_run_name = f"{algo}-multitask-summary" summary_log = _init_baseline_logger(cfg, summary_run_name) try: summary_payload: dict[str, float] = {} for key, stats in agg.items(): summary_payload[f"summary/{key}/mean"] = stats["mean"] summary_payload[f"summary/{key}/std"] = stats["std"] summary_log.log_summary(summary_payload) finally: summary_log.finish() logger.info("All %d seed(s) complete.", len(seeds))