Upload folder using huggingface_hub

examples/train_torch_ppo.py

@@ -173,11 +173,18 @@ def build_action_mask(observation: PyreObservation, exclude_look: bool = True) -
         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)
@@ -568,6 +575,22 @@ def run_episode(
                     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)
@@ -875,6 +898,29 @@ def build_curriculum(schedule_str: str, n_episodes: int) -> List[str]:
     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.
 
@@ -888,7 +934,13 @@ class PatienceCurriculum:
         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)
+        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__(
@@ -897,14 +949,28 @@ class PatienceCurriculum:
         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]
@@ -913,7 +979,7 @@ class PatienceCurriculum:
         """Call once per episode *after* appending to success_window.
 
         Returns the difficulty to use for the *next* episode.
-        Also handles the hard-phase medium-mix injection.
+        Also handles the final-stage cumulative-replay mix.
         """
         if self.stage_idx < len(self.stages) - 1:
             if success_rate_30 >= self.threshold:
@@ -929,11 +995,18 @@ class PatienceCurriculum:
                     f"for {self.patience_window} eps)"
                 )
 
-        # Hard-phase mix: occasionally replay medium to prevent forgetting
-        if self.current == "hard" and self.mix_ratio > 0.0 and len(self.stages) >= 2:
+        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  # medium replay episode
+                return prev
         return self.current
 
 
@@ -1019,13 +1092,17 @@ def train(args: argparse.Namespace) -> None:
     # 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:
@@ -1204,7 +1281,11 @@ def parse_args() -> argparse.Namespace:
                    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.")
+                        "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)

examples/train_torch_ppo_http.py

@@ -14,32 +14,31 @@ Usage
 2.  Run this script:
         .venv/Scripts/python.exe examples/train_torch_ppo_http.py
 
-Optional flags (identical to train_torch_ppo.py):
-    --server          Base URL of the Pyre server  [default: http://localhost:8000]
-    --episodes        Total training episodes       [default: 400]
-    --difficulty-schedule  Curriculum              [default: easy,easy,easy,medium,medium]
-    --output          Where to save the model .pt  [default: artifacts/pyre_ppo_http.pt]
-    ...               (all other flags are the same as train_torch_ppo.py)
+Optional flags (all match train_torch_ppo.py):
+    --server               Base URL of the Pyre server   [default: http://localhost:8000]
+    --episodes             Total training episodes        [default: 400]
+    --difficulty-schedule  Curriculum stages              [default: easy,medium,hard]
+    --patience-threshold   Success-rate gate (0=static)  [default: 0.65]
+    --learning-rate        Adam learning rate             [default: 3e-4]
+    --resume               Path to checkpoint to resume  [default: None]
+    --output               Where to save the model .pt   [default: artifacts/pyre_ppo_http.pt]
 """
 
 from __future__ import annotations
 
 import argparse
 import csv
-import os
 import sys
 import time
 from collections import deque
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 import numpy as np
 import requests
 import torch
-import torch.nn as nn
 import torch.optim as optim
-from torch.distributions import Categorical
 
 # ---------------------------------------------------------------------------
 # Resolve project root so we can import shared models regardless of CWD
@@ -67,10 +66,13 @@ from examples.train_torch_ppo import (
     WINDS,
     ActorCritic,
     ObservationEncoder,
+    PatienceCurriculum,
     RolloutBuffer,
     action_index_to_env_action,
     build_action_mask,
+    build_curriculum,
     compute_gae,
+    parse_mix_dist,
     ppo_update,
     save_training_graph_png,
 )
@@ -189,6 +191,7 @@ def run_episode(
     history_length: int,
     buffer: RolloutBuffer,
     deterministic: bool = False,
+    step_delay: float = 0.0,
 ) -> EpisodeResult:
     observation = env.reset(difficulty=difficulty)
     zero_frame = np.zeros(encoder.base_dim, dtype=np.float32)
@@ -215,6 +218,9 @@ def run_episode(
 
             action_idx = int(action_t.item())
             env_action = action_index_to_env_action(action_idx)
+            if step_delay > 0.0:
+                time.sleep(step_delay)
+            print(f"    step={steps+1:03d}  action={ACTION_KEYS[action_idx]:<40}  hp={observation.agent_health:5.1f}", flush=True)
             next_obs = env.step(env_action)
 
             reward = float(next_obs.reward or 0.0)
@@ -243,6 +249,17 @@ def run_episode(
                     reward -= 0.2
                 recent_positions.append(cur_pos)
 
+            # Shaping 4 — exit proximity pull
+            # Absolute distance-based bonus (not just delta) so the network
+            # has a continuous gradient toward exits from anywhere on the map.
+            # Max +0.25 when adjacent, tapers to 0 beyond 6 cells (Manhattan).
+            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
+                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)
@@ -278,9 +295,12 @@ def run_episode(
 
 def train(args: argparse.Namespace) -> None:
     device = torch.device("cuda" if args.device == "cuda" and torch.cuda.is_available() else "cpu")
+    if args.device == "cuda" and not torch.cuda.is_available():
+        print("[warn] CUDA not available — falling back to CPU.")
+
     encoder = ObservationEncoder(mode=args.observation_mode)
     input_dim = encoder.base_dim * args.history_length
-    hidden_sizes = [int(x) for x in args.hidden_sizes.split(",")]
+    hidden_sizes = tuple(int(x) for x in args.hidden_sizes.split(","))
     action_dim = ACTION_DIM
 
     # Connect to server
@@ -291,57 +311,114 @@ def train(args: argparse.Namespace) -> None:
         sys.exit(1)
     print("OK")
 
-    # Network
+    # Network + optimizer
     network = ActorCritic(input_dim, action_dim, hidden_sizes).to(device)
-    optimizer = optim.Adam(network.parameters(), lr=args.lr)
+    optimizer = optim.Adam(network.parameters(), lr=args.learning_rate, eps=1e-5)
+
+    # LinearLR scheduler: step once per PPO update, not per episode
+    total_updates = args.episodes // args.update_every
+    lr_scheduler = optim.lr_scheduler.LinearLR(
+        optimizer,
+        start_factor=1.0,
+        end_factor=args.lr_end_factor,
+        total_iters=max(1, total_updates),
+    ) if args.lr_decay else None
 
     total_params = sum(p.numel() for p in network.parameters())
     print(f"\n[config] server={args.server}")
     print(f"[config] device={device}  episodes={args.episodes}  batch={args.update_every} eps")
     print(f"[config] curriculum: {args.difficulty_schedule}")
-    print(f"[config] PPO clip_eps={args.clip_eps}  entropy={args.entropy_coef}  lr={args.lr}")
+    print(f"[config] PPO clip_eps={args.clip_eps}  entropy={args.entropy_coef}  lr={args.learning_rate}")
     print(f"\n[network] Parameters: {total_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)\n", flush=True)
 
-    schedule = args.difficulty_schedule.split(",")
+    # 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_schedule: Optional[List[str]] = None
+        print(f"[curriculum] patience-gated: threshold={args.patience_threshold}  "
+              f"window={args.patience_window}  mix={args.hard_mix_ratio}", flush=True)
+        if mix_dist is not None:
+            print(f"[curriculum] hard-phase mix distribution: {mix_dist}", flush=True)
+    else:
+        patience_curriculum = None
+        static_schedule = build_curriculum(args.difficulty_schedule, args.episodes)
+        print(f"[curriculum] static: {args.difficulty_schedule}", flush=True)
+
+    # Resume
+    start_ep = 0
+    if args.resume and Path(args.resume).exists():
+        ckpt = torch.load(args.resume, map_location="cpu", weights_only=False)
+        network.load_state_dict(ckpt.get("network_state", ckpt))
+        if "optimizer_state" in ckpt:
+            optimizer.load_state_dict(ckpt["optimizer_state"])
+        if lr_scheduler and ckpt.get("scheduler_state"):
+            lr_scheduler.load_state_dict(ckpt["scheduler_state"])
+        start_ep = int(ckpt.get("episode", 0))
+        print(f"[resume] Loaded checkpoint from episode {start_ep}: {args.resume}")
+
     buffer = RolloutBuffer()
-    metrics: list = []
-    eval_metrics: list = []
+    episode_rows: List[Dict] = []
+    eval_rows: List[Dict] = []
     success_window: deque = deque(maxlen=30)
     reward_window: deque = deque(maxlen=30)
     t0 = time.time()
-    lr_scheduler = optim.lr_scheduler.LinearLR(
-        optimizer, start_factor=1.0, end_factor=0.1, total_iters=args.episodes
-    )
 
-    for ep in range(1, args.episodes + 1):
-        stage_idx = min(int((ep - 1) / args.episodes * len(schedule)), len(schedule) - 1)
-        difficulty = schedule[stage_idx]
+    for ep_idx in range(start_ep, args.episodes):
+        ep = ep_idx + 1
+
+        # Determine difficulty for this episode
+        if patience_curriculum is not None:
+            difficulty = patience_curriculum.current
+        else:
+            difficulty = static_schedule[ep_idx]  # type: ignore[index]
+
+        # Use step delay only after --viz-after-ep episodes have been trained
+        ep_step_delay = args.step_delay if ep > args.viz_after_ep else 0.0
+        result = run_episode(env, network, encoder, device, difficulty, args.history_length, buffer,
+                             step_delay=ep_step_delay)
 
-        result = run_episode(env, network, encoder, device, difficulty, args.history_length, buffer)
         success_window.append(1 if result.evacuated else 0)
         reward_window.append(result.total_reward)
         suc30 = sum(success_window) / len(success_window)
         r30 = sum(reward_window) / len(reward_window)
         elapsed = int(time.time() - t0)
 
-        evac_sym = "1" if result.evacuated else "0"
+        # Advance patience curriculum after updating success_window
+        if patience_curriculum is not None:
+            difficulty = patience_curriculum.step(suc30)
+
         print(
             f"ep={ep:04d} [{difficulty:<6}] steps={result.steps:03d}  "
-            f"reward={result.total_reward:+8.3f}  evac={evac_sym}  "
+            f"reward={result.total_reward:+8.3f}  evac={int(result.evacuated)}  "
             f"hp={result.final_health:5.1f}  suc30={suc30:.2f}  "
-            f"r30={r30:+7.2f}  t={elapsed}s"
+            f"r30={r30:+7.2f}  t={elapsed}s",
+            flush=True,
         )
 
-        metrics.append({
-            "episode": ep, "difficulty": difficulty, "steps": result.steps,
-            "reward": round(result.total_reward, 4), "evacuated": int(result.evacuated),
-            "final_health": result.final_health, "suc30": round(suc30, 3), "r30": round(r30, 3),
+        episode_rows.append({
+            "episode": ep,
+            "difficulty": difficulty,
+            "steps": result.steps,
+            "reward": round(result.total_reward, 4),
+            "evacuated": int(result.evacuated),
+            "final_health": round(result.final_health, 2),
+            "reward_mean_30": round(r30, 4),
+            "success_rate_30": round(suc30, 4),
         })
 
-        # PPO update
-        if ep % args.update_every == 0 and len(buffer.obs) > 0:
+        # PPO update every N episodes (or at the very last episode)
+        should_update = (ep % args.update_every == 0) or (ep == args.episodes)
+        if should_update and len(buffer) > 0:
             network.train()
             stats = ppo_update(
                 network=network, optimizer=optimizer, buffer=buffer, device=device,
@@ -351,79 +428,98 @@ def train(args: argparse.Namespace) -> None:
                 gamma=args.gamma, gae_lambda=args.gae_lambda,
                 max_grad_norm=args.max_grad_norm,
             )
-            lr_scheduler.step()
+            if lr_scheduler:
+                lr_scheduler.step()
+            buffer.clear()
+            network.eval()
+
             cur_lr = optimizer.param_groups[0]["lr"]
             print(
                 f"  >> PPO update  samples=flushed  "
                 f"pi_loss={stats['policy_loss']:+.4f}  v_loss={stats['value_loss']:.4f}  "
                 f"entropy={stats['entropy']:.4f}  kl={stats['approx_kl']:.4f}  "
-                f"clip%={stats['clip_frac']:.2f}  lr={cur_lr:.2e}"
+                f"clip%={stats['clip_frac']:.2f}  lr={cur_lr:.2e}",
+                flush=True,
             )
-            buffer.clear()
-            network.eval()
 
         # Evaluation
-        if ep % args.eval_every == 0:
+        if args.eval_every > 0 and (ep % args.eval_every == 0 or ep == args.episodes):
             eval_rewards, eval_success, eval_steps_list = [], [], []
             eval_buf = RolloutBuffer()
             for _ in range(args.eval_episodes):
                 er = run_episode(
                     env, network, encoder, device,
                     args.eval_difficulty, args.history_length,
-                    eval_buf, deterministic=True,
+                    eval_buf, deterministic=True, step_delay=0.0,
                 )
+                eval_buf.clear()  # clear after each eval episode — don't accumulate
                 eval_rewards.append(er.total_reward)
                 eval_success.append(1 if er.evacuated else 0)
                 eval_steps_list.append(er.steps)
             avg_r = sum(eval_rewards) / len(eval_rewards)
             avg_s = sum(eval_success) / len(eval_success)
             avg_st = sum(eval_steps_list) / len(eval_steps_list)
-            print(f"  ** EVAL [{args.eval_difficulty}]  reward={avg_r:+.3f}  success={avg_s:.2f}  steps={avg_st:.1f}")
-            eval_metrics.append({
-                "episode": ep, "eval_difficulty": args.eval_difficulty,
-                "avg_reward": round(avg_r, 4), "success_rate": round(avg_s, 3),
-                "avg_steps": round(avg_st, 1),
+            print(f"  ** EVAL [{args.eval_difficulty}]  reward={avg_r:+.3f}  success={avg_s:.2f}  steps={avg_st:.1f}", flush=True)
+            eval_rows.append({
+                "episode": ep,
+                "difficulty": args.eval_difficulty,
+                "reward_mean": round(avg_r, 4),
+                "success_rate": round(avg_s, 3),
+                "steps_mean": round(avg_st, 1),
             })
 
-        # Checkpoint
-        if args.checkpoint and ep % args.checkpoint_every == 0:
-            torch.save(network.state_dict(), args.checkpoint)
-            print(f"  [ckpt] saved -> {args.checkpoint}")
+        # Periodic checkpoint (full state, same as train_torch_ppo.py)
+        if args.checkpoint and args.checkpoint_every > 0 and ep % args.checkpoint_every == 0:
+            ckpt_path = Path(args.checkpoint)
+            ckpt_path.parent.mkdir(parents=True, exist_ok=True)
+            torch.save({
+                "episode": ep,
+                "network_state": network.state_dict(),
+                "optimizer_state": optimizer.state_dict(),
+                "scheduler_state": lr_scheduler.state_dict() if lr_scheduler else None,
+                "args": vars(args),
+            }, ckpt_path)
+            print(f"  [ckpt] saved -> {args.checkpoint}", flush=True)
 
     # --- Save artefacts ---
     out = Path(args.output)
     out.parent.mkdir(parents=True, exist_ok=True)
-    torch.save(network.state_dict(), out)
+    torch.save({
+        "episode": args.episodes,
+        "network_state": network.state_dict(),
+        "optimizer_state": optimizer.state_dict(),
+        "scheduler_state": lr_scheduler.state_dict() if lr_scheduler else None,
+        "args": vars(args),
+    }, out)
     print(f"\n[done] Model saved -> {out}")
 
-    if args.save_metrics and metrics:
+    if args.save_metrics and episode_rows:
         csv_path = out.with_suffix(".csv")
         with open(csv_path, "w", newline="") as f:
-            writer = csv.DictWriter(f, fieldnames=metrics[0].keys())
+            writer = csv.DictWriter(f, fieldnames=episode_rows[0].keys())
             writer.writeheader()
-            writer.writerows(metrics)
+            writer.writerows(episode_rows)
         print(f"[done] Metrics CSV  -> {csv_path}")
 
-        if eval_metrics:
-            eval_csv = out.with_stem(out.stem + "_eval").with_suffix(".csv")
+        if eval_rows:
+            eval_csv = out.parent / (out.stem + "_eval.csv")
             with open(eval_csv, "w", newline="") as f:
-                writer = csv.DictWriter(f, fieldnames=eval_metrics[0].keys())
+                writer = csv.DictWriter(f, fieldnames=eval_rows[0].keys())
                 writer.writeheader()
-                writer.writerows(eval_metrics)
+                writer.writerows(eval_rows)
             print(f"[done] Eval CSV     -> {eval_csv}")
 
     if args.save_graph:
-        try:
-            png_path = out.with_suffix(".png")
-            save_training_graph_png(metrics, eval_metrics, str(png_path))
-            print(f"[done] Graph PNG    -> {png_path}")
-        except Exception as e:
-            print(f"[warn] Graph skipped: {e}")
+        png_path = out.with_suffix(".png")
+        # Correct arg order: save_training_graph_png(path, episode_rows, eval_rows)
+        save_training_graph_png(png_path, episode_rows, eval_rows)
+        print(f"[done] Graph PNG    -> {png_path}")
 
     suc_final = sum(success_window) / max(1, len(success_window))
     r_final = sum(reward_window) / max(1, len(reward_window))
     elapsed_total = time.time() - t0
-    print(f"\n[summary] {args.episodes} episodes in {elapsed_total:.1f}s  ({args.episodes / elapsed_total:.1f} eps/s)")
+    n_trained = args.episodes - start_ep
+    print(f"\n[summary] {n_trained} episodes in {elapsed_total:.1f}s  ({n_trained / max(1, elapsed_total):.1f} eps/s)")
     print(f"[summary] Final success rate (last 30): {suc_final:.2f}")
     print(f"[summary] Final reward mean  (last 30): {r_final:+.3f}")
 
@@ -434,19 +530,38 @@ def train(args: argparse.Namespace) -> None:
 
 def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(
-        description="PPO trainer using the Pyre HTTP server (localhost:8000)"
+        description="PPO trainer using the Pyre HTTP server (localhost:8000)",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
 
     # Server
     p.add_argument("--server", type=str, default="http://localhost:8000",
                    help="Base URL of the running Pyre env server")
 
-    # Training
+    # Visualization / pacing
+    p.add_argument("--step-delay", type=float, default=0.0,
+                   help="Seconds to sleep between steps (0=full speed, 0.5=smooth viz)")
+    p.add_argument("--viz-after-ep", type=int, default=0,
+                   help="Episode after which --step-delay activates. "
+                        "0=always delay, 100=fast first 100 eps then slow.")
+
+    # Training scale
     p.add_argument("--episodes", type=int, default=400)
+    p.add_argument("--max-steps", type=int, default=150, help="Max steps per episode (informational; enforced server-side)")
     p.add_argument("--device", type=str, default="cpu", choices=("cuda", "cpu"))
 
     # Curriculum
-    p.add_argument("--difficulty-schedule", type=str, default="easy,easy,easy,medium,medium")
+    p.add_argument("--difficulty-schedule", type=str, default="easy,medium,hard",
+                   help="Comma-separated curriculum stages")
+    p.add_argument("--patience-threshold", type=float, default=0.65,
+                   help="Success-rate (30-ep window) required before advancing difficulty. Set 0 for static split.")
+    p.add_argument("--patience-window", type=int, default=15,
+                   help="Consecutive episodes that must meet --patience-threshold before advancing.")
+    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("--hard-mix-ratio", type=float, default=0.25,
+                   help="Fraction of hard-phase episodes replayed on medium (prevents forgetting).")
     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)
@@ -456,10 +571,14 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--history-length", type=int, default=4)
 
     # Network
-    p.add_argument("--hidden-sizes", type=str, default="256,128,64")
-
-    # PPO
-    p.add_argument("--lr", type=float, default=3e-4)
+    p.add_argument("--hidden-sizes", type=str, default="512,256,128",
+                   help="Comma-separated MLP hidden layer sizes (match train_torch_ppo.py defaults)")
+
+    # PPO hyperparameters
+    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 updates")
+    p.add_argument("--lr-end-factor", type=float, default=0.1)
     p.add_argument("--gamma", type=float, default=0.99)
     p.add_argument("--gae-lambda", type=float, default=0.95)
     p.add_argument("--clip-eps", type=float, default=0.2)
@@ -470,10 +589,12 @@ def parse_args() -> argparse.Namespace:
     p.add_argument("--minibatch-size", type=int, default=256)
     p.add_argument("--max-grad-norm", type=float, default=0.5)
 
-    # Output
+    # Persistence
     p.add_argument("--output", type=str, default="artifacts/pyre_ppo_http.pt")
     p.add_argument("--checkpoint", type=str, default="artifacts/pyre_ppo_http_ckpt.pt")
     p.add_argument("--checkpoint-every", type=int, default=50)
+    p.add_argument("--resume", type=str, default=None,
+                   help="Path to a checkpoint (.pt) to resume training from")
     p.add_argument("--save-metrics", action="store_true", default=True)
     p.add_argument("--save-graph", action="store_true", default=True)
     p.add_argument("--seed", type=int, default=42)

server/fire_sim.py

@@ -34,9 +34,6 @@ FIRE_BURNING = 0.3
 FIRE_INTENSITY_GAIN = 0.15
 BURNOUT_TICKS = 5
 
-# Doors smolder longer before collapsing so they reliably spread fire into corridors
-DOOR_BURNOUT_MULTIPLIER = 3
-
 # Door fire reduction factor
 DOOR_CLOSED_FIRE_FACTOR = 0.15
 
@@ -221,14 +218,7 @@ class FireSim:
                     new_fire[i] = min(1.0, fire_grid[i] + intensity_gain)
                     if fire_grid[i] >= FIRE_BURNING:
                         new_burn_timers[i] = burn_timers[i] + 1
-                    # Doors smolder 3× longer than normal cells so they spread fire
-                    # into the corridor before collapsing to floor.
-                    effective_burnout = (
-                        self.burnout_ticks * DOOR_BURNOUT_MULTIPLIER
-                        if ct in (DOOR_OPEN, DOOR_CLOSED)
-                        else self.burnout_ticks
-                    )
-                    if new_burn_timers[i] >= effective_burnout and new_fire[i] >= 1.0:
+                    if new_burn_timers[i] >= self.burnout_ticks and new_fire[i] >= 1.0:
                         # Doors burn away to open floor (passage stays navigable).
                         # Everything else (furniture, walls that caught fire, etc.) becomes rubble.
                         cell_grid[i] = FLOOR if ct in (DOOR_OPEN, DOOR_CLOSED) else OBSTACLE

server/floor_plan.py

@@ -235,8 +235,8 @@ def _make_small_office() -> FloorPlan:
 #   Row 11: W F F O O F F F F F O O F F F W
 #   Row 12: W F F O O F F F F F O O F F F W
 #   Row 13: W F F F F F F F F F F F F F F W
-#   Row 14: W F F F F F F F F F F F F F F E  ← exit at (15,14)
-#   Row 15: W W W W W W W W W W W W W W W W
+#   Row 14: W F F F F F F F F F F F F F F W
+#   Row 15: W W W W W W W W W W W W W W W E  ← exit at (15,15)
 # ---------------------------------------------------------------------------
 
 def _make_open_plan() -> FloorPlan:
@@ -256,12 +256,12 @@ def _make_open_plan() -> FloorPlan:
         [1, 0, 0, 5, 5, 0, 0, 0, 0, 0, 5, 5, 0, 0, 0, 1],  # 11 pillars
         [1, 0, 0, 5, 5, 0, 0, 0, 0, 0, 5, 5, 0, 0, 0, 1],  # 12
         [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],  # 13
-        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4],  # 14 exit at x=15
-        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],  # 15
+        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1],  # 14
+        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 4],  # 15 exit at x=15
     ]
     grid = [c for row in rows for c in row]
 
-    exit_positions = [(0, 1), (15, 14)]
+    exit_positions = [(0, 1), (15, 15)]
     door_positions = []  # No internal doors in open plan
 
     floor_cells = [(x, y) for y in range(H) for x in range(W)
@@ -304,15 +304,14 @@ def _make_open_plan() -> FloorPlan:
 # Template 3: t_corridor
 #
 # T-shaped layout: vertical stem (x=7, y=0-14) + horizontal bar (y=7, x=0-15)
-# Side rooms off horizontal bar — accessed through bar-to-room doors at row 8,
-# with internal doors at row 10 connecting upper (y=8-9) and lower (y=11-12) halves.
+# Side rooms off horizontal bar (y=8-12, left and right of stem):
 #
 #   Row  0: W W W W W W W E W W W W W W W W  ← exit at (7,0)
 #   Row 1-6: vertical stem only (x=7)
 #   Row  7: E F F F F F F F F F F F F F F E  ← horizontal bar + exits
-#   Row  8: W W D W W D W F W W D W W W D W  ← bar-to-room doors (fire must pass these)
-#   Row  9: W F F W F F W F W F F W F F F W  ← room interior (4 cells deep: rows 9–12)
-#   Row 10: W F F W F F W F W F F W F F F W
+#   Row  8: W F F W F F W F W F F W F F F W  ← rooms branch off bar
+#   Row  9: W F F W F F W F W F F W F F F W
+#   Row 10: W W D W W D W F W D W W W W D W  ← doors to stem
 #   Row 11: W F F W F F W F W F F W F F F W
 #   Row 12: W F F W F F W F W F F W F F F W
 #   Row 13-14: stem only
@@ -330,11 +329,11 @@ def _make_t_corridor() -> FloorPlan:
         [1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1],  # 5
         [1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1],  # 6
         [4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4],  # 7  horizontal + exits
-        [1, 1, 2, 1, 1, 2, 1, 0, 1, 1, 2, 1, 1, 1, 2, 1],  # 8  bar→room doors
-        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 9  room interior
-        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 10 room interior
-        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 11 room interior
-        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 12 room interior (deepest)
+        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 8  side rooms
+        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 9
+        [1, 1, 2, 1, 1, 2, 1, 0, 1, 2, 1, 1, 1, 1, 2, 1],  # 10 doors to stem
+        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 11
+        [1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1],  # 12
         [1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1],  # 13 stem continues
         [1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1],  # 14
         [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],  # 15
@@ -342,11 +341,11 @@ def _make_t_corridor() -> FloorPlan:
     grid = [c for row in rows for c in row]
 
     exit_positions = [(7, 0), (0, 7), (15, 7)]
-    door_positions = [(2, 8), (5, 8), (10, 8), (14, 8)]  # bar → room entry doors
+    door_positions = [(2, 10), (5, 10), (9, 10), (14, 10)]
 
-    # Spawn zones: horizontal bar + side rooms (row 8 is now doors/walls, start from row 9)
+    # Spawn zones: horizontal bar + side rooms
     bar_cells = [(x, 7) for x in range(1, 15) if grid[_idx(x, 7, W)] == 0]
-    room_cells = [(x, y) for y in range(9, 13) for x in range(1, 15)
+    room_cells = [(x, y) for y in range(8, 13) for x in range(1, 15)
                   if grid[_idx(x, y, W)] == 0]
     stem_cells = [(7, y) for y in range(1, 15) if grid[_idx(7, y, W)] == 0]
 

server/pyre_env_environment.py

@@ -364,9 +364,8 @@ class PyreEnvironment(Environment):
         self._visited_cells.add((st.agent_x, st.agent_y))
 
         # Track closest approach to any exit for NearMissBonus
-        _exits = unblocked_exits(st.exit_positions, st.fire_grid, st.grid_w)
-        if not _exits:
-            _exits = st.exit_positions
+        _exits_reachable = unblocked_exits(st.exit_positions, st.fire_grid, st.grid_w)
+        _exits = _exits_reachable if _exits_reachable else st.exit_positions
         _cur_dist = bfs_exit_dist(st.agent_x, st.agent_y, _exits, st.cell_grid, st.grid_w, st.grid_h)
         if _cur_dist < self._min_exit_dist_reached:
             self._min_exit_dist_reached = _cur_dist
@@ -407,6 +406,12 @@ class PyreEnvironment(Environment):
         )
         obs_data["done"] = done
         obs_data["reward"] = reward
+
+        # Compute visible cells once so both metadata and map_state can use the count
+        _visible_set = compute_visible_cells(
+            st.agent_x, st.agent_y, st.cell_grid, st.smoke_grid, st.grid_w, st.grid_h,
+        ) if st.agent_alive and not st.agent_evacuated else set()
+
         obs_data["metadata"] = {
             "agent_health": st.agent_health,
             "step": st.step_count,
@@ -415,8 +420,12 @@ class PyreEnvironment(Environment):
             "fire_sources": st.fire_sources_count,
             "humidity": st.humidity,
             "difficulty": getattr(self, "_difficulty", "medium"),
+            # Fields consumed by ObservationEncoder global_features — previously missing
+            "nearest_exit_distance": _cur_dist,
+            "reachable_exit_count": len(_exits_reachable),
+            "visible_cell_count": len(_visible_set),
         }
-        obs_data["map_state"] = self._build_map_state(st)
+        obs_data["map_state"] = self._build_map_state(st, visible_set=_visible_set)
         return PyreObservation(**obs_data)
 
     @property
@@ -588,15 +597,22 @@ class PyreEnvironment(Environment):
     # Map state builder
     # ------------------------------------------------------------------
 
-    def _build_map_state(self, st: PyreState) -> PyreMapState:
-        """Assemble the full numerical grid snapshot for UI / visualization."""
+    def _build_map_state(self, st: PyreState, visible_set: Optional[set] = None) -> PyreMapState:
+        """Assemble the full numerical grid snapshot for UI / visualization.
+
+        Args:
+            visible_set: Pre-computed set of (x, y) visible cells. When provided
+                         (e.g. from step()) the second compute_visible_cells call
+                         is skipped. Pass None to compute fresh (used by reset()).
+        """
         if st.agent_alive and not st.agent_evacuated:
-            visible = compute_visible_cells(
-                st.agent_x, st.agent_y,
-                st.cell_grid, st.smoke_grid,
-                st.grid_w, st.grid_h,
-            )
-            visible_cells = [[x, y] for x, y in sorted(visible)]
+            if visible_set is None:
+                visible_set = compute_visible_cells(
+                    st.agent_x, st.agent_y,
+                    st.cell_grid, st.smoke_grid,
+                    st.grid_w, st.grid_h,
+                )
+            visible_cells = [[x, y] for x, y in sorted(visible_set)]
         else:
             visible_cells = []
 

server/rubrics.py

@@ -5,9 +5,9 @@ Each rubric class exposes a score() method.
 The environment composes them by calling each rubric each step.
 
 Per-step rubrics:
-  TimeStepPenalty            -0.05      constant time pressure
-  ProgressReward             +0.1       agent moved closer to nearest unblocked exit (BFS distance)
-  ProgressRegressionPenalty  -0.05      agent moved farther from nearest exit (symmetric gradient)
+  TimeStepPenalty            -0.01      constant time pressure
+  ProgressReward             +0.25      agent moved closer to nearest unblocked exit (BFS distance)
+  ProgressRegressionPenalty  -0.15      agent moved farther from nearest exit (symmetric gradient)
   SafeProgressBonus          +0.05      stacks on ProgressReward when progress made through smoke-free cell
   DangerPenalty              -0.5       agent moved into smoke≥moderate or fire-adjacent cell
   HealthDrainPenalty         -0.02×dmg  proportional to health lost this step
@@ -103,10 +103,10 @@ _bfs_exit_dist = bfs_exit_dist
 # ---------------------------------------------------------------------------
 
 class TimeStepPenalty:
-    """Constant penalty per step to encourage urgency."""
+    """Small constant penalty per step to encourage urgency."""
 
     def score(self, **_) -> float:
-        return -0.05
+        return -0.01
 
 
 class ProgressReward:
@@ -114,6 +114,8 @@ class ProgressReward:
 
     Uses BFS traversal distance (respects walls and obstacles) instead of
     Manhattan distance, so only genuine navigational progress is rewarded.
+    Value raised to +0.25 to create a stronger pull toward exits relative to
+    the danger/loop penalties that push the agent away from threats.
     """
 
     def score(
@@ -134,7 +136,7 @@ class ProgressReward:
             exits = exit_positions  # all blocked — still try to reward progress
         prev_dist = _bfs_exit_dist(prev_agent_x, prev_agent_y, exits, cell_grid, w, h)
         new_dist = _bfs_exit_dist(agent_x, agent_y, exits, cell_grid, w, h)
-        return 0.15 if new_dist < prev_dist else 0.0
+        return 0.25 if new_dist < prev_dist else 0.0
 
 
 class DangerPenalty:
@@ -219,8 +221,9 @@ class StrategicDoorBonus:
 class ProgressRegressionPenalty:
     """Penalise moving farther from the nearest unblocked exit.
 
-    Symmetric counterpart to ProgressReward: the agent gets +0.1 for progress
-    and –0.05 for regression, creating a clean two-sided gradient.
+    Symmetric counterpart to ProgressReward: the agent gets +0.25 for progress
+    and –0.15 for regression, creating a strong two-sided gradient that
+    discourages wandering away from the exit under fire pressure.
     """
 
     def score(
@@ -241,7 +244,7 @@ class ProgressRegressionPenalty:
             exits = exit_positions
         prev_dist = bfs_exit_dist(prev_agent_x, prev_agent_y, exits, cell_grid, w, h)
         new_dist = bfs_exit_dist(agent_x, agent_y, exits, cell_grid, w, h)
-        return -0.05 if new_dist > prev_dist else 0.0
+        return -0.15 if new_dist > prev_dist else 0.0
 
 
 class SafeProgressBonus: