examples/train_torch_ppo.py

"""
PyTorch PPO Agent for Pyre — Fire Evacuation RL Training Script.

=== ENVIRONMENT SUMMARY ===
Pyre is a partial-observability crisis navigation environment:
  - Grid: 16×16 (easy/medium) or 20×24 (hard, procedural)
  - Agent: Spawns inside a burning building, must evacuate before dying
  - Fire: Spreads via cellular automaton — wind, humidity, fuel vary per episode
  - Partial observability: visibility radius (2–5 cells) shrinks in heavy smoke
  - Doors: Can be opened/closed to slow fire spread (+0.5 strategic door bonus)
  - Health: 100 HP, drains from smoke (0.5–5/step) and fire (10/step)

=== ACTION SPACE (41 discrete) ===
  0–3   : move(north|south|west|east)
  4–7   : look(north|south|west|east)  — scan without moving, still costs a step
  8     : wait()
  9–24  : door(door_1..16, open)
  25–40 : door(door_1..16, close)
  Runtime action masking via `available_actions_hint` prevents invalid moves.

=== OBSERVATION ENCODING ===
  Per-step grid: 24×24 padded map × 10 channels
    • 6 one-hot cell type (floor/wall/door_open/door_closed/exit/obstacle)
    • fire intensity [0, 1]
    • smoke density  [0, 1]
    • visibility mask (1=visible, 0=unseen)
    • agent position mask
  Global scalars (22): health, step_progress, fire_spread, humidity,
    agent_x, agent_y, exit_distance, reachable_exits, visible_cells,
    fire_sources, smoke_severity, alive, evacuated, wind (one-hot 5), difficulty (one-hot 3)
  Frame stacking: 4 consecutive frames → input_dim = 5782 × 4 = 23128

=== REWARD STRUCTURE ===
  Per-step:
    -0.01  time penalty (urgency)
    +0.10  BFS progress toward nearest unblocked exit
    -0.05  regression (moved farther from exit)
    +0.05  safe-progress bonus (progress through smoke-free cell)
    -0.50  danger penalty (moved into smoke≥moderate or fire-adjacent)
    -0.02×dmg health drain penalty
    +0.50  strategic door close (adjacent to fire, once per door per episode)
    +0.02  exploration bonus (first visit to cell)
  Terminal:
    +5.00  evacuation success
    +1.50×(hp/100) health survival bonus (max +1.5)
    -10.0  death
    -5.00  timeout
    0→+3.0 near-miss partial credit (based on closest exit approach)
    +0.05×remaining_steps time bonus

=== ALGORITHM: PPO (Proximal Policy Optimization) ===
WHY PPO over alternatives:
  • DQN    — Off-policy, harder credit assignment for sparse terminal rewards; no clean action masking
  • A2C    — Simpler but no clipping → unstable on hard stochastic episodes
  • SAC    — Designed for continuous spaces; discrete SAC works but adds complexity
  • LSTM-PPO — Better for fully text-only obs; grid map_state already encodes spatial state
  → PPO + frame-stack + action-mask hits the sweet spot for this env

Key PPO improvements over the existing NumPy A2C (train_rl_agent.py):
  ✓ PPO clip (ε=0.2)        prevents catastrophic updates
  ✓ Entropy regularization  sustains exploration in smoke-obscured corridors
  ✓ Value function clipping  stabilises critic under sparse terminal rewards
  ✓ GPU acceleration         10–20× faster than NumPy baseline
  ✓ LayerNorm in network     improves gradient flow for large input dims
  ✓ Linear LR decay          stabilises late-stage convergence
  ✓ Better curriculum        3-stage easy→medium→hard with patience gating

Usage:
    python examples/train_torch_ppo.py --episodes 500 --device cuda
    python examples/train_torch_ppo.py --episodes 300 --difficulty-schedule easy,medium,hard
    python examples/train_torch_ppo.py --resume artifacts/pyre_ppo_checkpoint.pt
    python examples/train_torch_ppo.py --describe-only
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np

# ---------------------------------------------------------------------------
# Optional torch import — fail fast with a helpful message
# ---------------------------------------------------------------------------
try:
    import torch
    import torch.nn as nn
    import torch.nn.functional as F
    from torch.optim import Adam
    from torch.optim.lr_scheduler import LinearLR
except ImportError:
    sys.exit(
        "PyTorch not found. Install with:\n"
        "  pip install torch --index-url https://download.pytorch.org/whl/cu121\n"
        "or for CPU only:\n"
        "  pip install torch"
    )

# ---------------------------------------------------------------------------
# Project imports — support both package install and direct run from root
# ---------------------------------------------------------------------------
_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

try:
    from pyre_env.models import PyreAction, PyreObservation
    from pyre_env.server.pyre_env_environment import PyreEnvironment
except ModuleNotFoundError:
    try:
        from models import PyreAction, PyreObservation
        from server.pyre_env_environment import PyreEnvironment
    except ModuleNotFoundError:
        sys.exit(
            "Cannot import Pyre modules. Run this script from the openenv-pyre root:\n"
            "  python examples/train_torch_ppo.py"
        )

# ---------------------------------------------------------------------------
# Reuse the established observation/action interface from train_rl_agent.py
# These are the canonical definitions for this environment.
# ---------------------------------------------------------------------------
MAX_GRID_W = 24
MAX_GRID_H = 24
MAX_DOORS = 16
DIRECTIONS = ("north", "south", "west", "east")
WINDS = ("CALM", "NORTH", "SOUTH", "WEST", "EAST")
DIFFICULTIES = ("easy", "medium", "hard")

MOVE_KEYS = [f"move(direction='{d}')" for d in DIRECTIONS]
LOOK_KEYS = [f"look(direction='{d}')" for d in DIRECTIONS]
WAIT_KEY = "wait()"
OPEN_KEYS = [f"door(target_id='door_{i}', door_state='open')" for i in range(1, MAX_DOORS + 1)]
CLOSE_KEYS = [f"door(target_id='door_{i}', door_state='close')" for i in range(1, MAX_DOORS + 1)]
ACTION_KEYS = MOVE_KEYS + LOOK_KEYS + [WAIT_KEY] + OPEN_KEYS + CLOSE_KEYS
ACTION_DIM = len(ACTION_KEYS)  # 41
ACTION_TO_INDEX = {key: idx for idx, key in enumerate(ACTION_KEYS)}

import re
_MOVE_RE = re.compile(r"move\(direction='(north|south|west|east)'\)")
_LOOK_RE = re.compile(r"look\(direction='(north|south|west|east)'\)")
_DOOR_RE = re.compile(r"door\(target_id='(door_(\d+))', door_state='(open|close)'\)")


def action_index_to_env_action(index: int) -> PyreAction:
    if 0 <= index < 4:
        return PyreAction(action="move", direction=DIRECTIONS[index])
    if 4 <= index < 8:
        return PyreAction(action="look", direction=DIRECTIONS[index - 4])
    if index == 8:
        return PyreAction(action="wait")
    if 9 <= index < 9 + MAX_DOORS:
        door_id = f"door_{index - 8}"
        return PyreAction(action="door", target_id=door_id, door_state="open")
    door_slot = index - (9 + MAX_DOORS)
    door_id = f"door_{door_slot + 1}"
    return PyreAction(action="door", target_id=door_id, door_state="close")


def build_action_mask(observation: PyreObservation, exclude_look: bool = True) -> np.ndarray:
    """Build a binary validity mask over the 41-action space.

    exclude_look=True (default for RL):
        Suppresses all 4 'look' actions. The RL agent already receives the full
        grid via map_state — look gives zero new information but wastes a step
        and earns no reward. Excluding it concentrates the policy on moves and
        doors, which are the only actions that can improve the agent's position.

    NOTE: Look action indices are 4–7 in ACTION_KEYS. The guard below must be
    applied in the ACTION_TO_INDEX fast-path as well as the regex fallback,
    because look hint strings exactly match ACTION_TO_INDEX keys and would
    otherwise bypass the exclude_look flag entirely.
    """
    mask = np.zeros(ACTION_DIM, dtype=np.float32)
    for hint in observation.available_actions_hint:
        idx = ACTION_TO_INDEX.get(hint)
        if idx is not None:
            if exclude_look and 4 <= idx <= 7:  # indices 4-7 are look(north/south/west/east)
                continue
            mask[idx] = 1.0
            continue
        m = _MOVE_RE.fullmatch(hint)
        if m:
            mask[ACTION_TO_INDEX[f"move(direction='{m.group(1)}')"]] = 1.0
            continue
        m = _LOOK_RE.fullmatch(hint)
        if m:
            if not exclude_look:
                mask[ACTION_TO_INDEX[f"look(direction='{m.group(1)}')"]] = 1.0
            continue
        m = _DOOR_RE.fullmatch(hint)
        if m:
            door_id, door_num, state = m.group(1), int(m.group(2)), m.group(3)
            if 1 <= door_num <= MAX_DOORS:
                mask[ACTION_TO_INDEX[f"door(target_id='{door_id}', door_state='{state}')"]] = 1.0
    if mask.sum() == 0:
        mask[ACTION_TO_INDEX[WAIT_KEY]] = 1.0
    return mask


class ObservationEncoder:
    """Encode PyreObservation into a fixed-length float32 vector.

    Mode 'visible': only populate cells within the agent's sight radius —
        mimics true partial observability; preferred for training.
    Mode 'full': expose complete ground-truth grid — useful for debugging
        or oracle upper-bound experiments.

    Output shape: (base_dim,) = (MAX_GRID_W × MAX_GRID_H × 10 + 25,) = (5785,)
    With history stacking of k frames: (5785 × k,)

    The 3 extra scalars over the v1 baseline are map-agnostic exit-compass
    features (Fix 3): exit_dx_norm, exit_dy_norm, exit_manhattan_norm.
    These allow the agent to locate the nearest exit on procedurally generated
    maps without having to memorise layout-specific coordinates.
    """

    base_dim = MAX_GRID_W * MAX_GRID_H * 10 + 25

    def __init__(self, mode: str = "visible"):
        if mode not in {"visible", "full"}:
            raise ValueError(f"mode must be 'visible' or 'full', got '{mode}'")
        self.mode = mode

    def encode(self, observation: PyreObservation) -> np.ndarray:
        ms = observation.map_state
        if ms is None:
            raise ValueError("map_state is required for encoding.")

        cell_one_hot = np.zeros((MAX_GRID_H, MAX_GRID_W, 6), dtype=np.float32)
        fire_ch = np.zeros((MAX_GRID_H, MAX_GRID_W), dtype=np.float32)
        smoke_ch = np.zeros((MAX_GRID_H, MAX_GRID_W), dtype=np.float32)
        vis_ch = np.zeros((MAX_GRID_H, MAX_GRID_W), dtype=np.float32)
        agent_ch = np.zeros((MAX_GRID_H, MAX_GRID_W), dtype=np.float32)

        visible = {(x, y) for x, y in ms.visible_cells}
        for y in range(ms.grid_h):
            for x in range(ms.grid_w):
                if self.mode == "visible" and (x, y) not in visible and (x, y) != (ms.agent_x, ms.agent_y):
                    continue
                i = y * ms.grid_w + x
                ct = int(ms.cell_grid[i])
                if 0 <= ct <= 5:
                    cell_one_hot[y, x, ct] = 1.0
                fire_ch[y, x] = float(ms.fire_grid[i])
                smoke_ch[y, x] = float(ms.smoke_grid[i])
                vis_ch[y, x] = 1.0 if (x, y) in visible else 0.0

        if 0 <= ms.agent_x < MAX_GRID_W and 0 <= ms.agent_y < MAX_GRID_H:
            agent_ch[ms.agent_y, ms.agent_x] = 1.0

        grid_features = np.concatenate([
            cell_one_hot.reshape(-1),
            fire_ch.reshape(-1),
            smoke_ch.reshape(-1),
            vis_ch.reshape(-1),
            agent_ch.reshape(-1),
        ])

        meta = observation.metadata or {}
        wind = str(meta.get("wind_dir", ms.wind_dir or "CALM")).upper()
        diff = str(meta.get("difficulty", "medium")).lower()
        wi = WINDS.index(wind) if wind in WINDS else 0
        di = DIFFICULTIES.index(diff) if diff in DIFFICULTIES else 1

        wind_oh = np.zeros(len(WINDS), dtype=np.float32); wind_oh[wi] = 1.0
        diff_oh = np.zeros(len(DIFFICULTIES), dtype=np.float32); diff_oh[di] = 1.0

        # Fix 3 — map-agnostic exit compass features.
        # Compute the direction vector and normalised Manhattan distance to the
        # nearest exit cell (cell_type == 4) directly from the live grid.
        # This gives the agent an exit "compass" that works on procedurally
        # generated maps without memorising any layout.
        EXIT_CELL_TYPE = 4
        ax, ay = ms.agent_x, ms.agent_y
        gw, gh = ms.grid_w, ms.grid_h
        best_dist = float(gw + gh)
        best_dx = 0.0
        best_dy = 0.0
        for cy in range(gh):
            for cx in range(gw):
                if int(ms.cell_grid[cy * gw + cx]) == EXIT_CELL_TYPE:
                    d = abs(cx - ax) + abs(cy - ay)
                    if d < best_dist:
                        best_dist = d
                        best_dx = float(cx - ax) / max(1, gw - 1)
                        best_dy = float(cy - ay) / max(1, gh - 1)
        exit_manhattan_norm = best_dist / float(gw + gh)

        global_features = np.array([
            float(observation.agent_health) / 100.0,
            float(ms.agent_health) / 100.0,
            float(ms.step_count) / max(1, ms.max_steps),
            float(ms.fire_spread_rate),
            float(ms.humidity),
            float(ms.agent_x) / max(1, ms.grid_w - 1),
            float(ms.agent_y) / max(1, ms.grid_h - 1),
            float(meta.get("nearest_exit_distance", MAX_GRID_W + MAX_GRID_H) or 0.0) / float(MAX_GRID_W + MAX_GRID_H),
            float(meta.get("reachable_exit_count", 0.0)) / 4.0,
            float(meta.get("visible_cell_count", 0.0)) / float(MAX_GRID_W * MAX_GRID_H),
            float(meta.get("fire_sources", 0.0)) / 5.0,
            {"none": 0.0, "light": 0.33, "moderate": 0.66, "heavy": 1.0}.get(observation.smoke_level, 0.0),
            1.0 if ms.agent_alive else 0.0,
            1.0 if ms.agent_evacuated else 0.0,
            # Fix 3: exit-compass (3 new scalars — map-agnostic, layout-independent)
            best_dx,           # signed x-direction toward nearest exit
            best_dy,           # signed y-direction toward nearest exit
            exit_manhattan_norm,  # how far away the exit is (0 = here, 1 = max)
        ], dtype=np.float32)

        return np.concatenate([grid_features, global_features, wind_oh, diff_oh]).astype(np.float32)


# ---------------------------------------------------------------------------
# Neural Network
# ---------------------------------------------------------------------------

class ActorCritic(nn.Module):
    """Shared-backbone Actor-Critic network for PPO.

    Architecture:
        Input → LayerNorm → FC(512) → LayerNorm → ReLU
                          → FC(256) → LayerNorm → ReLU
                          → FC(128) → ReLU
               ┌──────────────┴──────────────┐
         Policy head (→ logits)        Value head (→ scalar)

    LayerNorm before activations improves gradient flow for the large
    (23128-dim) flat input without requiring feature normalization.
    """

    def __init__(self, input_dim: int, action_dim: int, hidden_sizes: Tuple[int, ...] = (512, 256, 128)):
        super().__init__()
        h1, h2, h3 = hidden_sizes

        self.shared = nn.Sequential(
            nn.LayerNorm(input_dim),
            nn.Linear(input_dim, h1),
            nn.LayerNorm(h1),
            nn.ReLU(),
            nn.Linear(h1, h2),
            nn.LayerNorm(h2),
            nn.ReLU(),
            nn.Linear(h2, h3),
            nn.ReLU(),
        )

        # Orthogonal init — standard for PPO (improves early convergence)
        self._init_orthogonal()

        self.policy_head = nn.Linear(h3, action_dim)
        self.value_head = nn.Linear(h3, 1)

        # Small init for output heads prevents saturated softmax early on
        nn.init.orthogonal_(self.policy_head.weight, gain=0.01)
        nn.init.zeros_(self.policy_head.bias)
        nn.init.orthogonal_(self.value_head.weight, gain=1.0)
        nn.init.zeros_(self.value_head.bias)

    def _init_orthogonal(self) -> None:
        for layer in self.shared:
            if isinstance(layer, nn.Linear):
                nn.init.orthogonal_(layer.weight, gain=np.sqrt(2))
                nn.init.zeros_(layer.bias)

    def forward(
        self,
        obs: torch.Tensor,
        mask: torch.Tensor,
    ) -> Tuple[torch.distributions.Categorical, torch.Tensor]:
        """
        Args:
            obs:  (B, input_dim) float32
            mask: (B, action_dim) float32  — 1.0 = valid, 0.0 = invalid
        Returns:
            dist:   Categorical distribution (action masking applied as -inf)
            values: (B,) float32
        """
        features = self.shared(obs)
        logits = self.policy_head(features)

        # Mask invalid actions with -inf before softmax (numerically stable)
        logits = torch.where(mask.bool(), logits, torch.full_like(logits, -1e9))

        dist = torch.distributions.Categorical(logits=logits)
        values = self.value_head(features).squeeze(-1)
        return dist, values

    def act(
        self,
        obs: torch.Tensor,
        mask: torch.Tensor,
        deterministic: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Sample (or take greedy) action. Returns (action, log_prob, value)."""
        dist, values = self(obs, mask)
        action = dist.mode if deterministic else dist.sample()
        log_prob = dist.log_prob(action)
        return action, log_prob, values

    def evaluate(
        self,
        obs: torch.Tensor,
        mask: torch.Tensor,
        action: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Evaluate stored actions during PPO update. Returns (log_prob, value, entropy)."""
        dist, values = self(obs, mask)
        log_prob = dist.log_prob(action)
        entropy = dist.entropy()
        return log_prob, values, entropy


# ---------------------------------------------------------------------------
# Rollout buffer
# ---------------------------------------------------------------------------

@dataclass
class RolloutBuffer:
    """Stores transitions for a batch of episodes before PPO update."""
    obs: List[np.ndarray] = field(default_factory=list)
    masks: List[np.ndarray] = field(default_factory=list)
    actions: List[int] = field(default_factory=list)
    rewards: List[float] = field(default_factory=list)
    log_probs: List[float] = field(default_factory=list)
    values: List[float] = field(default_factory=list)
    dones: List[bool] = field(default_factory=list)

    def clear(self) -> None:
        self.obs.clear()
        self.masks.clear()
        self.actions.clear()
        self.rewards.clear()
        self.log_probs.clear()
        self.values.clear()
        self.dones.clear()

    def __len__(self) -> int:
        return len(self.rewards)


# ---------------------------------------------------------------------------
# GAE computation
# ---------------------------------------------------------------------------

def compute_gae(
    rewards: np.ndarray,
    values: np.ndarray,
    dones: np.ndarray,
    gamma: float,
    gae_lambda: float,
) -> Tuple[np.ndarray, np.ndarray]:
    """Generalized Advantage Estimation.

    Returns (returns, advantages) — both shape (T,).
    Episode boundaries (done=True) reset the GAE accumulator so advantages
    don't bleed across episodes within a mixed batch.
    """
    T = len(rewards)
    advantages = np.zeros(T, dtype=np.float32)
    gae = 0.0
    next_value = 0.0
    for t in reversed(range(T)):
        if dones[t]:
            next_value = 0.0
            gae = 0.0
        delta = rewards[t] + gamma * next_value * (1.0 - dones[t]) - values[t]
        gae = delta + gamma * gae_lambda * (1.0 - dones[t]) * gae
        advantages[t] = gae
        next_value = values[t]
    returns = advantages + values
    return returns, advantages


# ---------------------------------------------------------------------------
# Episode runner
# ---------------------------------------------------------------------------

@dataclass
class EpisodeResult:
    total_reward: float
    steps: int
    evacuated: bool
    final_health: float
    difficulty: str


def run_episode(
    env: PyreEnvironment,
    network: ActorCritic,
    encoder: ObservationEncoder,
    device: torch.device,
    difficulty: str,
    history_length: int,
    buffer: RolloutBuffer,
    deterministic: bool = False,
) -> EpisodeResult:
    """Run one episode, appending transitions to *buffer*."""
    observation = env.reset(difficulty=difficulty)
    zero_frame = np.zeros(encoder.base_dim, dtype=np.float32)
    frames: deque = deque([zero_frame.copy() for _ in range(history_length)], maxlen=history_length)
    frames.append(encoder.encode(observation))

    total_reward = 0.0
    final_health = observation.agent_health
    evacuated = False
    steps = 0
    # Anti-loop tracking: remember the last LOOP_WINDOW positions this episode.
    # Revisiting any of them means the agent is circling, not exploring.
    LOOP_WINDOW = 12
    recent_positions: deque = deque(maxlen=LOOP_WINDOW)

    network.eval()
    with torch.no_grad():
        while True:
            state_vec = np.concatenate(list(frames), dtype=np.float32)
            # exclude_look=True: RL agent sees full grid — look wastes steps
            action_mask = build_action_mask(observation, exclude_look=True)

            obs_t = torch.tensor(state_vec, dtype=torch.float32, device=device).unsqueeze(0)
            mask_t = torch.tensor(action_mask, dtype=torch.float32, device=device).unsqueeze(0)

            action_t, log_prob_t, value_t = network.act(obs_t, mask_t, deterministic=deterministic)

            action_idx = int(action_t.item())
            env_action = action_index_to_env_action(action_idx)
            next_obs = env.step(env_action)

            reward = float(next_obs.reward or 0.0)

            # ----------------------------------------------------------------
            # Reward shaping 1 — idle penalty
            # The env's -0.01/step is too weak; make waiting explicitly costly.
            # ----------------------------------------------------------------
            chosen_action = env_action.action
            if chosen_action == "wait":
                reward -= 0.05

            # ----------------------------------------------------------------
            # Reward shaping 2 — fire-approach penalty (Fix 2)
            # Penalise landing on (or moving next to) a cell with active fire.
            # This is stronger than the env's DangerPenalty and fires *before*
            # health drain accumulates, teaching the agent to predict spread.
            # We look at the NEW observation's map to catch the current step.
            # ----------------------------------------------------------------
            ms_next = next_obs.map_state
            if ms_next is not None and chosen_action.startswith("move"):
                ax, ay = ms_next.agent_x, ms_next.agent_y
                gw, gh = ms_next.grid_w, ms_next.grid_h
                fire_grid = ms_next.fire_grid
                for dx, dy in ((0, 1), (0, -1), (1, 0), (-1, 0)):
                    nx, ny = ax + dx, ay + dy
                    if 0 <= nx < gw and 0 <= ny < gh:
                        if float(fire_grid[ny * gw + nx]) > 0.15:
                            reward -= 0.15  # early fire-proximity warning
                            break

            # ----------------------------------------------------------------
            # Reward shaping 3 — anti-loop penalty
            # If the agent steps onto a cell it occupied in the last LOOP_WINDOW
            # steps, it is circling. Penalise to force forward exploration.
            # Fires only on move actions — wait is already penalised above.
            # ----------------------------------------------------------------
            if ms_next is not None and chosen_action.startswith("move"):
                cur_pos = (ms_next.agent_x, ms_next.agent_y)
                if cur_pos in recent_positions:
                    reward -= 0.2  # break the loop
                recent_positions.append(cur_pos)

            # ----------------------------------------------------------------
            # Reward shaping 4 — exit proximity pull
            # Absolute (not just delta) distance-based bonus so the agent has
            # a continuous gradient toward exits even before it learns
            # consistent BFS progress.  Complements the server-side
            # ProgressReward which only fires on a single step of BFS gain.
            # Max +0.25 when adjacent; tapers to 0 beyond 6 cells (Manhattan).
            # Only fires on move to avoid rewarding standing still near exits.
            # ----------------------------------------------------------------
            if ms_next is not None and chosen_action.startswith("move") and not next_obs.agent_evacuated:
                ax, ay = ms_next.agent_x, ms_next.agent_y
                exits = ms_next.exit_positions  # List[List[int]] of [x, y]
                if exits:
                    min_manhattan = min(abs(ax - ex[0]) + abs(ay - ex[1]) for ex in exits)
                    reward += max(0.0, 0.25 - 0.04 * min_manhattan)

            done = bool(next_obs.done)

            buffer.obs.append(state_vec)
            buffer.masks.append(action_mask)
            buffer.actions.append(action_idx)
            buffer.rewards.append(reward)
            buffer.log_probs.append(float(log_prob_t.item()))
            buffer.values.append(float(value_t.item()))
            buffer.dones.append(done)

            total_reward += reward
            steps += 1
            final_health = next_obs.agent_health
            evacuated = next_obs.agent_evacuated

            frames.append(encoder.encode(next_obs))
            observation = next_obs
            if done:
                break

    return EpisodeResult(
        total_reward=total_reward,
        steps=steps,
        evacuated=evacuated,
        final_health=final_health,
        difficulty=difficulty,
    )


# ---------------------------------------------------------------------------
# PPO update
# ---------------------------------------------------------------------------

def ppo_update(
    network: ActorCritic,
    optimizer: Adam,
    buffer: RolloutBuffer,
    device: torch.device,
    clip_eps: float,
    value_clip_eps: float,
    entropy_coef: float,
    value_coef: float,
    n_epochs: int,
    minibatch_size: int,
    gamma: float,
    gae_lambda: float,
    max_grad_norm: float,
) -> Dict[str, float]:
    """Full PPO update over the collected rollout buffer."""
    rewards = np.array(buffer.rewards, dtype=np.float32)
    values = np.array(buffer.values, dtype=np.float32)
    dones = np.array(buffer.dones, dtype=np.float32)

    returns, advantages = compute_gae(rewards, values, dones, gamma, gae_lambda)

    # Normalize advantages across the whole batch (reduces variance)
    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

    obs_arr = torch.tensor(np.stack(buffer.obs), dtype=torch.float32, device=device)
    mask_arr = torch.tensor(np.stack(buffer.masks), dtype=torch.float32, device=device)
    action_arr = torch.tensor(buffer.actions, dtype=torch.long, device=device)
    old_logp_arr = torch.tensor(buffer.log_probs, dtype=torch.float32, device=device)
    return_arr = torch.tensor(returns, dtype=torch.float32, device=device)
    adv_arr = torch.tensor(advantages, dtype=torch.float32, device=device)
    old_value_arr = torch.tensor(values, dtype=torch.float32, device=device)

    T = len(buffer)
    metrics = {"policy_loss": 0.0, "value_loss": 0.0, "entropy": 0.0, "approx_kl": 0.0, "clip_frac": 0.0}
    n_updates = 0

    network.train()
    for _ in range(n_epochs):
        perm = torch.randperm(T, device=device)
        for start in range(0, T, minibatch_size):
            idx = perm[start:start + minibatch_size]
            if len(idx) < 2:
                continue

            log_prob, value, entropy = network.evaluate(obs_arr[idx], mask_arr[idx], action_arr[idx])

            # PPO ratio and clipped surrogate loss
            ratio = torch.exp(log_prob - old_logp_arr[idx])
            adv_mb = adv_arr[idx]
            surr1 = ratio * adv_mb
            surr2 = torch.clamp(ratio, 1.0 - clip_eps, 1.0 + clip_eps) * adv_mb
            policy_loss = -torch.min(surr1, surr2).mean()

            # Value loss with optional clipping (stabilises critic)
            ret_mb = return_arr[idx]
            old_val_mb = old_value_arr[idx]
            value_pred_clipped = old_val_mb + torch.clamp(value - old_val_mb, -value_clip_eps, value_clip_eps)
            value_loss = torch.max(
                F.mse_loss(value, ret_mb),
                F.mse_loss(value_pred_clipped, ret_mb),
            )

            entropy_loss = -entropy.mean()

            loss = policy_loss + value_coef * value_loss + entropy_coef * entropy_loss

            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(network.parameters(), max_grad_norm)
            optimizer.step()

            with torch.no_grad():
                approx_kl = ((ratio - 1) - (log_prob - old_logp_arr[idx])).mean().item()
                clip_frac = ((ratio - 1.0).abs() > clip_eps).float().mean().item()

            metrics["policy_loss"] += policy_loss.item()
            metrics["value_loss"] += value_loss.item()
            metrics["entropy"] += entropy.mean().item()
            metrics["approx_kl"] += approx_kl
            metrics["clip_frac"] += clip_frac
            n_updates += 1

    if n_updates > 0:
        for k in metrics:
            metrics[k] /= n_updates
    return metrics


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

def evaluate_policy(
    env: PyreEnvironment,
    network: ActorCritic,
    encoder: ObservationEncoder,
    device: torch.device,
    difficulty: str,
    history_length: int,
    n_episodes: int,
) -> Dict[str, float]:
    rewards, successes, steps = [], [], []
    dummy_buffer = RolloutBuffer()
    for _ in range(n_episodes):
        result = run_episode(
            env=env, network=network, encoder=encoder, device=device,
            difficulty=difficulty, history_length=history_length,
            buffer=dummy_buffer, deterministic=True,
        )
        dummy_buffer.clear()
        rewards.append(result.total_reward)
        successes.append(float(result.evacuated))
        steps.append(result.steps)
    return {
        "reward_mean": float(np.mean(rewards)),
        "reward_max": float(np.max(rewards)),
        "success_rate": float(np.mean(successes)),
        "steps_mean": float(np.mean(steps)),
    }


# ---------------------------------------------------------------------------
# PNG graph (matplotlib)
# ---------------------------------------------------------------------------

def save_training_graph_png(
    path: Path,
    episode_rows: List[Dict],
    eval_rows: List[Dict],
    window: int = 20,
) -> None:
    """Save a publication-quality PNG training graph with dual Y-axes."""
    try:
        import matplotlib
        matplotlib.use("Agg")   # non-interactive backend — no display needed
        import matplotlib.pyplot as plt
        import matplotlib.ticker as mticker
    except ImportError:
        print("[warn] matplotlib not installed — skipping PNG graph. Run: uv pip install matplotlib")
        return

    if not episode_rows:
        return

    path.parent.mkdir(parents=True, exist_ok=True)

    episodes   = [int(r["episode"]) for r in episode_rows]
    rewards    = [float(r["reward"]) for r in episode_rows]
    evacuated  = [float(r["evacuated"]) for r in episode_rows]
    difficulty = [str(r["difficulty"]) for r in episode_rows]

    # Moving average helper
    def ma(values: list, w: int) -> list:
        out, run, q = [], 0.0, []
        for v in values:
            q.append(v); run += v
            if len(q) > w: run -= q.pop(0)
            out.append(run / len(q))
        return out

    reward_ma  = ma(rewards, window)
    success_ma = ma(evacuated, window)

    eval_eps  = [int(r["episode"])      for r in eval_rows]
    eval_succ = [float(r["success_rate"]) for r in eval_rows]

    # Difficulty shading regions
    diff_colors = {"easy": "#d4edda", "medium": "#fff3cd", "hard": "#f8d7da"}
    regions: List[tuple] = []
    if difficulty:
        cur, start = difficulty[0], episodes[0]
        for ep, d in zip(episodes[1:], difficulty[1:]):
            if d != cur:
                regions.append((start, ep, cur))
                cur, start = d, ep
        regions.append((start, episodes[-1], cur))

    fig, ax1 = plt.subplots(figsize=(14, 6))
    ax2 = ax1.twinx()

    # Shade difficulty regions
    for x0, x1, diff in regions:
        ax1.axvspan(x0, x1, color=diff_colors.get(diff, "#eeeeee"), alpha=0.35, zorder=0)

    # Zero line
    ax1.axhline(0, color="#aaaaaa", linewidth=0.8, linestyle="--", zorder=1)

    # Raw reward (faint)
    ax1.plot(episodes, rewards, color="#d1c7bc", linewidth=0.8,
             alpha=0.6, label="Episode reward", zorder=2)

    # Reward moving average
    ax1.plot(episodes, reward_ma, color="#c1661c", linewidth=2.5,
             label=f"Reward (MA-{window})", zorder=3)

    # Success moving average (right axis)
    ax2.plot(episodes, success_ma, color="#1a7a8a", linewidth=2.5,
             linestyle="-", label=f"Success rate (MA-{window})", zorder=3)

    # Eval checkpoints
    if eval_eps:
        ax2.scatter(eval_eps, eval_succ, color="#0d5b6b", s=60, zorder=5,
                    marker="D", label="Eval success", edgecolors="white", linewidths=1.2)

    # Axes labels & formatting
    ax1.set_xlabel("Episode", fontsize=13, fontweight="bold", labelpad=8)
    ax1.set_ylabel("Reward", fontsize=13, fontweight="bold", color="#c1661c", labelpad=8)
    ax2.set_ylabel("Success Rate", fontsize=13, fontweight="bold", color="#1a7a8a", labelpad=8)

    ax1.tick_params(axis="y", labelcolor="#c1661c")
    ax2.tick_params(axis="y", labelcolor="#1a7a8a")
    ax2.set_ylim(-0.05, 1.05)
    ax2.yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1.0, decimals=0))

    ax1.grid(True, which="major", linestyle="--", linewidth=0.6,
             color="#dddddd", alpha=0.8, zorder=0)
    ax1.set_xlim(episodes[0], episodes[-1])

    ax1.tick_params(axis="x", labelsize=10)
    ax1.tick_params(axis="y", labelsize=10)
    ax2.tick_params(axis="y", labelsize=10)

    # Title
    total_eps = episodes[-1]
    final_sr  = success_ma[-1] if success_ma else 0.0
    fig.suptitle(
        f"Pyre PPO Training  —  {total_eps} episodes  |  final success rate: {final_sr:.0%}",
        fontsize=14, fontweight="bold", y=1.01,
    )

    # Difficulty legend patches
    import matplotlib.patches as mpatches
    diff_patches = [
        mpatches.Patch(color=diff_colors[d], alpha=0.6, label=d.capitalize())
        for d in ["easy", "medium", "hard"] if any(r == d for r in difficulty)
    ]

    # Combine legends from both axes
    h1, l1 = ax1.get_legend_handles_labels()
    h2, l2 = ax2.get_legend_handles_labels()
    ax1.legend(h1 + h2 + diff_patches, l1 + l2 + [p.get_label() for p in diff_patches],
               loc="upper left", fontsize=9, framealpha=0.85)

    fig.tight_layout()
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


# ---------------------------------------------------------------------------
# Curriculum scheduling
# ---------------------------------------------------------------------------

def build_curriculum(schedule_str: str, n_episodes: int) -> List[str]:
    """Expand comma-separated difficulty stages evenly over n_episodes.

    Example: 'easy,medium,hard' with 300 episodes → 100 each.
    Used only when patience_threshold=0 (static schedule).
    """
    stages = [s.strip().lower() for s in schedule_str.split(",") if s.strip()]
    if not stages:
        stages = ["medium"]
    for s in stages:
        if s not in DIFFICULTIES:
            raise ValueError(f"Unknown difficulty '{s}'. Choose from {DIFFICULTIES}.")
    seg = max(1, n_episodes // len(stages))
    schedule = []
    for s in stages:
        schedule.extend([s] * seg)
    while len(schedule) < n_episodes:
        schedule.append(stages[-1])
    return schedule[:n_episodes]


def parse_mix_dist(spec: Optional[str]) -> Optional[Dict[str, float]]:
    """Parse a 'hard:0.6,medium:0.3,easy:0.1' style spec into a dict.

    Returns None when ``spec`` is falsy. Probabilities are renormalised to
    sum to 1 if they don't already (within 1% tolerance).
    """
    if not spec:
        return None
    out: Dict[str, float] = {}
    for chunk in spec.split(","):
        chunk = chunk.strip()
        if not chunk:
            continue
        if ":" not in chunk:
            raise ValueError(f"Invalid mix-dist entry '{chunk}', expected 'name:prob'")
        name, val = chunk.split(":", 1)
        out[name.strip().lower()] = float(val)
    total = sum(out.values())
    if total <= 0:
        raise ValueError(f"mix-dist probabilities must be positive, got {out}")
    return {k: v / total for k, v in out.items()}


class PatienceCurriculum:
    """Dynamic difficulty scheduler that gates advancement on sustained success rate.

    Stays on current difficulty until success_rate_30 >= threshold for
    patience_window consecutive episodes, then advances to the next stage.
    During the hard phase an optional mix_ratio fraction of episodes are
    replayed on the previous (medium) difficulty to prevent catastrophic
    forgetting of the medium policy.

    Args:
        stages:           ordered list of difficulty strings, e.g. ['easy','medium','hard']
        threshold:        minimum success rate (0–1) required before advancing
        patience_window:  number of consecutive episodes that must meet threshold
        mix_ratio:        fraction of hard-phase episodes to run on medium instead (0–1).
                          Ignored when ``mix_dist`` is provided.
        mix_dist:         optional dict mapping difficulty -> probability used
                          during the *final* (hard) stage, e.g.
                          ``{"hard": 0.6, "medium": 0.3, "easy": 0.1}``. When set,
                          each hard-phase episode samples its difficulty from this
                          distribution. Probabilities must sum to 1.
    """

    def __init__(
        self,
        stages: List[str],
        threshold: float,
        patience_window: int,
        mix_ratio: float = 0.0,
        mix_dist: Optional[Dict[str, float]] = None,
    ) -> None:
        self.stages = stages
        self.threshold = threshold
        self.patience_window = patience_window
        self.mix_ratio = mix_ratio
        self.mix_dist = mix_dist
        self.stage_idx = 0
        self._streak = 0

        if self.mix_dist is not None:
            total = sum(self.mix_dist.values())
            if not (0.99 <= total <= 1.01):
                raise ValueError(
                    f"mix_dist probabilities must sum to 1, got {total:.3f}"
                )
            for k in self.mix_dist:
                if k not in self.stages:
                    raise ValueError(
                        f"mix_dist key '{k}' not in stages {self.stages}"
                    )

    @property
    def current(self) -> str:
        return self.stages[self.stage_idx]

    def step(self, success_rate_30: float) -> str:
        """Call once per episode *after* appending to success_window.

        Returns the difficulty to use for the *next* episode.
        Also handles the final-stage cumulative-replay mix.
        """
        if self.stage_idx < len(self.stages) - 1:
            if success_rate_30 >= self.threshold:
                self._streak += 1
            else:
                self._streak = 0
            if self._streak >= self.patience_window:
                self.stage_idx += 1
                self._streak = 0
                print(
                    f"  [curriculum] Advanced to '{self.current}' "
                    f"(success_rate_30={success_rate_30:.2f} >= {self.threshold} "
                    f"for {self.patience_window} eps)"
                )

        is_final_stage = self.stage_idx == len(self.stages) - 1

        if is_final_stage and self.mix_dist is not None:
            keys = list(self.mix_dist.keys())
            probs = np.array([self.mix_dist[k] for k in keys], dtype=np.float64)
            probs = probs / probs.sum()
            return str(np.random.choice(keys, p=probs))

        if is_final_stage and self.mix_ratio > 0.0 and len(self.stages) >= 2:
            prev = self.stages[self.stage_idx - 1]
            if np.random.rand() < self.mix_ratio:
                return prev
        return self.current


# ---------------------------------------------------------------------------
# Checkpoint
# ---------------------------------------------------------------------------

def save_checkpoint(
    path: Path,
    network: ActorCritic,
    optimizer: Adam,
    scheduler,
    episode: int,
    args: argparse.Namespace,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({
        "episode": episode,
        "network_state": network.state_dict(),
        "optimizer_state": optimizer.state_dict(),
        "scheduler_state": scheduler.state_dict() if scheduler else None,
        "args": vars(args),
    }, path)


def load_checkpoint(
    path: Path,
    network: ActorCritic,
    optimizer: Adam,
    scheduler,
) -> int:
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    network.load_state_dict(ckpt["network_state"])
    optimizer.load_state_dict(ckpt["optimizer_state"])
    if scheduler and ckpt.get("scheduler_state"):
        scheduler.load_state_dict(ckpt["scheduler_state"])
    start_episode = int(ckpt.get("episode", 0))
    print(f"[resume] Loaded checkpoint from episode {start_episode}: {path}")
    return start_episode


# ---------------------------------------------------------------------------
# CSV logging
# ---------------------------------------------------------------------------

def save_csv(path: Path, rows: List[Dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    if not rows:
        return
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()))
        writer.writeheader()
        writer.writerows(rows)


# ---------------------------------------------------------------------------
# Main training loop
# ---------------------------------------------------------------------------

def train(args: argparse.Namespace) -> None:
    device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
    if args.device == "cuda" and not torch.cuda.is_available():
        print("[warn] CUDA not available - falling back to CPU.")

    print(f"[config] device={device}  episodes={args.episodes}  batch={args.update_every} eps  "
          f"hidden={args.hidden_sizes}  frames={args.history_length}")
    print(f"[config] curriculum: {args.difficulty_schedule}")
    print(f"[config] PPO clip_eps={args.clip_eps}  entropy={args.entropy_coef}  lr={args.learning_rate}\n")

    encoder = ObservationEncoder(mode=args.observation_mode)
    input_dim = encoder.base_dim * args.history_length

    hidden_sizes = tuple(int(h) for h in args.hidden_sizes.split(","))
    network = ActorCritic(input_dim=input_dim, action_dim=ACTION_DIM, hidden_sizes=hidden_sizes).to(device)
    optimizer = Adam(network.parameters(), lr=args.learning_rate, eps=1e-5)

    total_steps_for_scheduler = args.episodes // args.update_every
    scheduler = LinearLR(optimizer, start_factor=1.0, end_factor=args.lr_end_factor,
                          total_iters=max(1, total_steps_for_scheduler)) if args.lr_decay else None

    env = PyreEnvironment(max_steps=args.max_steps)

    # Build curriculum — patience-gated (dynamic) or static
    stages = [s.strip().lower() for s in args.difficulty_schedule.split(",") if s.strip()]
    if args.patience_threshold > 0:
        mix_dist = parse_mix_dist(getattr(args, "hard_mix_dist", None))
        patience_curriculum = PatienceCurriculum(
            stages=stages,
            threshold=args.patience_threshold,
            patience_window=args.patience_window,
            mix_ratio=args.hard_mix_ratio,
            mix_dist=mix_dist,
        )
        static_curriculum: Optional[List[str]] = None
        if mix_dist is not None:
            print(f"[curriculum] hard-phase mix distribution: {mix_dist}")
        print(f"[curriculum] patience-gated: threshold={args.patience_threshold}  "
              f"window={args.patience_window}  mix={args.hard_mix_ratio}")
    else:
        patience_curriculum = None
        static_curriculum = build_curriculum(args.difficulty_schedule, args.episodes)
        print(f"[curriculum] static: {args.difficulty_schedule}")

    start_episode = 0
    if args.resume:
        resume_path = Path(args.resume)
        if resume_path.exists():
            start_episode = load_checkpoint(resume_path, network, optimizer, scheduler)

    # Tracking
    buffer = RolloutBuffer()
    episode_rows: List[Dict] = []
    eval_rows: List[Dict] = []
    reward_window: deque = deque(maxlen=30)
    success_window: deque = deque(maxlen=30)

    n_params = sum(p.numel() for p in network.parameters())
    print(f"[network] Parameters: {n_params:,}")
    print(f"[network] Input dim:  {input_dim:,}  (encoder.base_dim={encoder.base_dim} x {args.history_length} frames)")
    print(f"[network] Action dim: {ACTION_DIM}  (4 move + 4 look + 1 wait + {MAX_DOORS} open + {MAX_DOORS} close)")
    print()

    t_start = time.time()

    for ep_idx in range(start_episode, args.episodes):
        # Determine difficulty for this episode
        if patience_curriculum is not None:
            difficulty = patience_curriculum.current
        else:
            difficulty = static_curriculum[ep_idx]  # type: ignore[index]

        result = run_episode(
            env=env, network=network, encoder=encoder, device=device,
            difficulty=difficulty, history_length=args.history_length,
            buffer=buffer, deterministic=False,
        )

        reward_window.append(result.total_reward)
        success_window.append(float(result.evacuated))

        # Advance patience curriculum *after* updating success_window
        if patience_curriculum is not None:
            difficulty = patience_curriculum.step(float(np.mean(success_window)))

        ep_num = ep_idx + 1
        episode_rows.append({
            "episode": ep_num,
            "difficulty": difficulty,
            "reward": round(result.total_reward, 4),
            "evacuated": int(result.evacuated),
            "steps": result.steps,
            "final_health": round(result.final_health, 2),
            "reward_mean_30": round(float(np.mean(reward_window)), 4),
            "success_rate_30": round(float(np.mean(success_window)), 4),
        })

        elapsed = time.time() - t_start
        print(
            f"ep={ep_num:04d} [{difficulty:<6}] "
            f"steps={result.steps:03d}  "
            f"reward={result.total_reward:+8.3f}  "
            f"evac={int(result.evacuated)}  "
            f"hp={result.final_health:5.1f}  "
            f"suc30={float(np.mean(success_window)):.2f}  "
            f"r30={float(np.mean(reward_window)):+7.2f}  "
            f"t={elapsed:.0f}s"
        )

        # PPO update every N episodes
        should_update = (ep_num % args.update_every == 0) or (ep_num == args.episodes)
        if should_update and len(buffer) > 0:
            ppo_metrics = ppo_update(
                network=network, optimizer=optimizer, buffer=buffer, device=device,
                clip_eps=args.clip_eps, value_clip_eps=args.clip_eps,
                entropy_coef=args.entropy_coef, value_coef=args.value_coef,
                n_epochs=args.update_epochs, minibatch_size=args.minibatch_size,
                gamma=args.gamma, gae_lambda=args.gae_lambda,
                max_grad_norm=args.max_grad_norm,
            )
            if scheduler:
                scheduler.step()
            buffer.clear()

            cur_lr = optimizer.param_groups[0]["lr"]
            print(
                f"  >> PPO update  samples={len(buffer) if len(buffer) > 0 else 'flushed'}  "
                f"pi_loss={ppo_metrics['policy_loss']:+.4f}  "
                f"v_loss={ppo_metrics['value_loss']:.4f}  "
                f"entropy={ppo_metrics['entropy']:.4f}  "
                f"kl={ppo_metrics['approx_kl']:.4f}  "
                f"clip%={ppo_metrics['clip_frac']:.2f}  "
                f"lr={cur_lr:.2e}"
            )

        # Periodic evaluation
        if args.eval_every > 0 and (ep_num % args.eval_every == 0 or ep_num == args.episodes):
            eval_m = evaluate_policy(
                env=env, network=network, encoder=encoder, device=device,
                difficulty=args.eval_difficulty, history_length=args.history_length,
                n_episodes=args.eval_episodes,
            )
            eval_rows.append({"episode": ep_num, "difficulty": args.eval_difficulty, **{k: round(v, 4) for k, v in eval_m.items()}})
            print(
                f"  ** EVAL [{args.eval_difficulty}]  "
                f"reward={eval_m['reward_mean']:+.3f}  "
                f"success={eval_m['success_rate']:.2f}  "
                f"steps={eval_m['steps_mean']:.1f}"
            )

        # Periodic checkpoint
        if args.checkpoint and args.checkpoint_every > 0 and ep_num % args.checkpoint_every == 0:
            save_checkpoint(Path(args.checkpoint), network, optimizer, scheduler, ep_num, args)
            print(f"  [ckpt] saved -> {args.checkpoint}")

    # Final save
    if args.output:
        out = Path(args.output)
        save_checkpoint(out, network, optimizer, scheduler, args.episodes, args)
        print(f"\n[done] Model saved -> {out}")

        if args.save_metrics:
            csv_path = out.with_suffix(".csv")
            save_csv(csv_path, episode_rows)
            print(f"[done] Metrics CSV  -> {csv_path}")

        if eval_rows:
            eval_csv = out.parent / (out.stem + "_eval.csv")
            save_csv(eval_csv, eval_rows)
            print(f"[done] Eval CSV     -> {eval_csv}")

        if args.save_graph:
            png_path = out.with_suffix(".png")
            save_training_graph_png(png_path, episode_rows, eval_rows)
            print(f"[done] Graph PNG    -> {png_path}")

    total_time = time.time() - t_start
    print(f"\n[summary] {args.episodes - start_episode} episodes in {total_time:.1f}s  "
          f"({(args.episodes - start_episode) / max(1, total_time):.1f} eps/s)")
    print(f"[summary] Final success rate (last 30): {float(np.mean(success_window)):.2f}")
    print(f"[summary] Final reward mean  (last 30): {float(np.mean(reward_window)):+.3f}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def describe_env() -> None:
    print(__doc__)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="PPO training for Pyre fire-evacuation environment",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # Training scale
    p.add_argument("--episodes", type=int, default=400, help="Total training episodes")
    p.add_argument("--max-steps", type=int, default=150, help="Max steps per episode")
    p.add_argument("--device", type=str, default="cuda", choices=("cuda", "cpu"), help="Torch device")

    # Curriculum
    p.add_argument("--difficulty", type=str, default="easy", choices=DIFFICULTIES,
                   help="Single difficulty (overridden by --difficulty-schedule if set)")
    p.add_argument("--difficulty-schedule", type=str, default="easy,medium,hard",
                   help="Comma-separated curriculum stages. With --patience-threshold>0 these "
                        "become gated stages; otherwise split evenly across episodes.")
    p.add_argument("--patience-threshold", type=float, default=0.65,
                   help="Success-rate threshold (30-ep window) required before advancing to next "
                        "difficulty. Set 0 to use static even-split schedule.")
    p.add_argument("--patience-window", type=int, default=15,
                   help="Episodes that must sustain >= patience-threshold before advancing.")
    p.add_argument("--hard-mix-ratio", type=float, default=0.25,
                   help="Fraction of hard-phase episodes to replay on medium (0=pure hard). "
                        "Prevents catastrophic forgetting of the medium policy. "
                        "Ignored when --hard-mix-dist is set.")
    p.add_argument("--hard-mix-dist", type=str, default=None,
                   help="Cumulative replay distribution for the final stage, e.g. "
                        "'hard:0.6,medium:0.3,easy:0.1'. Overrides --hard-mix-ratio.")
    p.add_argument("--eval-difficulty", type=str, default="medium", choices=DIFFICULTIES)
    p.add_argument("--eval-episodes", type=int, default=10)
    p.add_argument("--eval-every", type=int, default=50)

    # Observation
    p.add_argument("--observation-mode", type=str, default="visible", choices=("visible", "full"),
                   help="'visible': partial obs (realistic); 'full': oracle grid (debug)")
    p.add_argument("--history-length", type=int, default=4,
                   help="Frames stacked per observation (temporal context for partial obs)")

    # Network
    p.add_argument("--hidden-sizes", type=str, default="512,256,128",
                   help="Comma-separated MLP hidden layer sizes")

    # PPO hyperparameters
    p.add_argument("--update-every", type=int, default=5,
                   help="Episodes between PPO updates (smaller = faster feedback loop early in training)")
    p.add_argument("--update-epochs", type=int, default=4,
                   help="Gradient passes over each collected batch (PPO allows >1)")
    p.add_argument("--minibatch-size", type=int, default=256)
    p.add_argument("--clip-eps", type=float, default=0.2, help="PPO surrogate clip ε")
    p.add_argument("--entropy-coef", type=float, default=0.03,
                   help="Entropy bonus coefficient — higher = more exploration (0.03 default encourages early exit-seeking)")
    p.add_argument("--value-coef", type=float, default=0.5)
    p.add_argument("--gamma", type=float, default=0.99)
    p.add_argument("--gae-lambda", type=float, default=0.95)
    p.add_argument("--max-grad-norm", type=float, default=0.5)

    # Optimizer / LR schedule
    p.add_argument("--learning-rate", type=float, default=3e-4)
    p.add_argument("--lr-decay", action="store_true", default=True,
                   help="Linear LR decay to lr_end_factor × initial_lr over training")
    p.add_argument("--lr-end-factor", type=float, default=0.1,
                   help="LR at end of training = initial_lr × this value")

    # Persistence
    p.add_argument("--output", type=str, default="artifacts/pyre_ppo.pt",
                   help="Path to save final model checkpoint")
    p.add_argument("--checkpoint", type=str, default="artifacts/pyre_ppo_checkpoint.pt",
                   help="Path for periodic checkpoints (also used by --resume)")
    p.add_argument("--checkpoint-every", type=int, default=50)
    p.add_argument("--resume", type=str, default=None,
                   help="Path to checkpoint to resume training from")
    p.add_argument("--save-metrics", action="store_true", default=True,
                   help="Save per-episode metrics as CSV alongside the model")
    p.add_argument("--save-graph", action="store_true", default=True,
                   help="Save a PNG training graph alongside the model (requires matplotlib)")

    # Misc
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--describe-only", action="store_true",
                   help="Print environment/algorithm description and exit")

    return p.parse_args()


def main() -> None:
    args = parse_args()

    if args.describe_only:
        describe_env()
        return

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    train(args)


if __name__ == "__main__":
    main()