Fix All: Cleaned All Files

Dockerfile

@@ -1,9 +1,3 @@
-# server/Dockerfile
-# ─────────────────────────────────────────────────────────────────────────────
-# Builds the Cascade Containment environment server.
-# Exposes port 7860 — required for Hugging Face Spaces deployment.
-# ─────────────────────────────────────────────────────────────────────────────
-
 FROM python:3.11-slim
 
 WORKDIR /app
@@ -23,4 +17,4 @@ ENV PYTHONPATH="/app:/app/server"
 
 EXPOSE 7860
 
-CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]

baseline/evaluator.py

@@ -1,4 +1,3 @@
-# baseline/evaluator.py
 import os, sys, time
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 from typing import List, Tuple, Any
@@ -13,20 +12,17 @@ from core.policy_update import compute_advantage, update_memory
 import requests as http_requests
 
 N_ROLLOUTS = {
-    "easy":   3,   # Always solves cleanly in 1-2 rollouts, no variance to learn from
-    "medium": 4,   # Needs GRPO learning signal to stabilise
-    "hard":   4,   # Needs GRPO learning signal to stabilise
+    "easy":   3,
+    "medium": 4,
+    "hard":   4,
 }
 
 
 def build_prompt_with_memory(obs: CityObservation, memory: EpisodicMemory) -> str:
-    from baseline.policy import build_prompt
     base         = build_prompt(obs)
     memory_block = memory.retrieve(obs)
-
     if not memory_block:
         return base
-
     injection = "\n" + memory_block + "\nApply these lessons to your current decision.\n"
     return base.replace("Your response:", injection + "Your response:")
 
@@ -129,8 +125,10 @@ def run_task_grpo(
         if verbose:
             mean = sum(completed_rewards[:-1]) / max(len(completed_rewards) - 1, 1) \
                    if len(completed_rewards) > 1 else total_reward
-            print(f"    → Advantage: {advantage:+.4f} | "
-                + (f"↑ Stored {stored} steps" if stored > 0 else "↓ Suppressed"))
+            print(
+                f"    → Advantage: {advantage:+.4f} | "
+                + (f"↑ Stored {stored} steps" if stored > 0 else "↓ Suppressed")
+            )
 
     all_rewards = [r[0] for r in rollouts]
     mean_reward = sum(all_rewards) / len(all_rewards)
@@ -188,4 +186,4 @@ def run_evaluation(base_url: str = "http://localhost:7860", verbose: bool = True
         print(f"  Time:    {elapsed}s")
         print("="*52 + "\n")
 
-    return scores
+    return scores

baseline/policy.py

@@ -1,4 +1,3 @@
-# baseline/policy.py
 import os, sys, json, re
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 from openai import OpenAI
@@ -51,7 +50,7 @@ def build_prompt(obs: CityObservation) -> str:
             hosp_status = "hospital OK"
 
         if has_data_lag:
-            # Pre-compute estimated current infection — don't ask LLM to do math
+            # pre-compute estimated current infection so the LLM doesn't have to do the math
             estimated = round(min(1.0, d.reported_infection_rate + 3 * d.growth_rate_hint), 2)
             if estimated > 0.4:
                 est_status = "🔴 EST.CRITICAL"
@@ -142,7 +141,7 @@ def call_llm(prompt: str, client: OpenAI) -> str:
             },
             {"role": "user", "content": prompt}
         ],
-        max_tokens  = 60,   # increased to allow brief reasoning + JSON
+        max_tokens  = 60,
         temperature = 0.1,
     )
     return (response.choices[0].message.content or "").strip()
@@ -169,4 +168,4 @@ def parse_action(response: str, num_districts: int) -> ContainmentAction:
 def get_action(obs: CityObservation, client: OpenAI) -> ContainmentAction:
     prompt   = build_prompt(obs)
     response = call_llm(prompt, client)
-    return parse_action(response, len(obs.districts))
+    return parse_action(response, len(obs.districts))

baseline/run.py

@@ -1,32 +1,24 @@
-# baseline/run.py
-# ─────────────────────────────────────────────────────────────────────────────
-# CLI entry point for the baseline evaluation.
-# Called by inference.py — can also be run directly for testing.
-# ─────────────────────────────────────────────────────────────────────────────
-
 import os
 import sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
 from dotenv import load_dotenv
 load_dotenv()
-print(f"DEBUG TOKEN: '{os.environ.get('HF_TOKEN', 'NOT SET')[:10]}...'")
 
-# If HF_TOKEN not set in .env, fall back to the HF CLI cache file
+# Fall back to the HuggingFace CLI token cache if HF_TOKEN isn't in .env
 if not os.environ.get("HF_TOKEN"):
     cache_path = os.path.expanduser("~/.cache/huggingface/token")
     if os.path.exists(cache_path):
         with open(cache_path, "r") as f:
             os.environ["HF_TOKEN"] = f.read().strip()
-        print(f"✓ Loaded HF_TOKEN from cache: {os.environ['HF_TOKEN'][:8]}...")
 
 from baseline.evaluator import run_evaluation
 
+
 def main():
     base_url = os.environ.get("ENV_BASE_URL", "http://localhost:7860")
-    scores   = run_evaluation(base_url=base_url, verbose=True)
-    return scores
+    return run_evaluation(base_url=base_url, verbose=True)
 
 
 if __name__ == "__main__":
-    main()
+    main()

client.py

@@ -1,11 +1,3 @@
-# client.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Client-side interface for the Cascade Containment environment.
-# Implements the two required abstract methods from EnvClient:
-#   _step_payload  — serialises ContainmentAction to dict for WebSocket
-#   _parse_result  — deserialises server response to CityObservation
-# ─────────────────────────────────────────────────────────────────────────────
-
 from openenv.core.env_client import EnvClient
 from openenv.core.client_types import StepResult
 from openenv.core.env_server.types import State
@@ -14,53 +6,49 @@ from models import ContainmentAction, CityObservation
 
 class CascadeContainmentEnv(EnvClient[ContainmentAction, CityObservation, State]):
     """
-    Client for the Cascade Containment OpenEnv environment.
+    Client for the Cascade Containment environment.
 
-    Async usage:
+    Async:
         async with CascadeContainmentEnv(base_url="http://localhost:7860") as env:
             obs = await env.reset("easy")
             result = await env.step(ContainmentAction(action_type="allocate", district_id=0))
 
-    Sync usage:
+    Sync:
         with CascadeContainmentEnv(base_url="http://localhost:7860").sync() as env:
             obs = env.reset("easy")
             result = env.step(ContainmentAction(action_type="allocate", district_id=0))
     """
 
     def _step_payload(self, action: ContainmentAction) -> dict:
-        """Serialise ContainmentAction to dict for WebSocket transmission."""
         return {
             "action_type": action.action_type,
             "district_id": action.district_id,
         }
 
     def _parse_result(self, result: dict) -> StepResult:
-        """Deserialise server response into a typed StepResult."""
         observation = CityObservation(**result["observation"])
         return StepResult(
             observation = observation,
             reward      = result.get("reward", 0.0),
             done        = result.get("done", False),
         )
-    
+
     def _parse_state(self, result: dict) -> State:
-        """Deserialise server response into a typed State."""
         return State(
-            episode_id  = result.get("episode_id", ""),
-            step_count  = result.get("step_count", 0),
+            episode_id = result.get("episode_id", ""),
+            step_count = result.get("step_count", 0),
         )
 
 
-# ── Connection test (run directly to verify client works) ─────────────────────
-
 if __name__ == "__main__":
+    # Quick smoke test — run directly to verify the client connects and steps correctly.
     with CascadeContainmentEnv(base_url="http://localhost:7860").sync() as env:
         obs = env.reset()
-        print(f"✓ Connected successfully")
-        print(f"  Districts:  {len(obs.observation.districts)}")
-        print(f"  Resources:  {obs.observation.available_resources}")
-        print(f"  Max steps:  {obs.observation.max_steps}")
+        print(f"✓ Connected")
+        print(f"  Districts:   {len(obs.observation.districts)}")
+        print(f"  Resources:   {obs.observation.available_resources}")
+        print(f"  Max steps:   {obs.observation.max_steps}")
 
         result = env.step(ContainmentAction(action_type="allocate", district_id=0))
         print(f"  Step reward: {result.reward}")
-        print(f"✓ Client working end-to-end")
+        print(f"✓ End-to-end OK")

core/policy_update.py

@@ -1,25 +1,19 @@
-# core/policy_update.py
-# ─────────────────────────────────────────────────────────────────────────────
-# GRPO-style advantage computation and memory update logic.
-# Determines which rollouts are above average and should be reinforced.
-# ─────────────────────────────────────────────────────────────────────────────
-
 import os
 import sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
-from typing import List, Tuple
+from typing import List
 from core.trajectory import EpisodicMemory
 
 
 def compute_advantage(
-    current_reward: float,
+    current_reward:    float,
     completed_rewards: List[float],
 ) -> float:
     """
-    GRPO advantage = R_i - mean(R).
-    Positive advantage → this rollout was better than average → reinforce.
-    Negative advantage → below average → suppress.
+    GRPO advantage = R_i - mean(R_completed).
+    Positive means this rollout was better than average; negative means worse.
+    Returns 0.0 on the first rollout where there's nothing to compare against.
     """
     if not completed_rewards:
         return 0.0
@@ -28,11 +22,8 @@ def compute_advantage(
 
 
 def should_reinforce(advantage: float) -> bool:
-    """
-    Reinforce if advantage >= 0 (at or above mean).
-    Suppress if below mean.
-    """
-    return advantage > -0.5  # allow small negative margin to encourage exploration
+    # Small negative margin is allowed to encourage exploration on borderline rollouts.
+    return advantage > -0.5
 
 
 def update_memory(
@@ -41,10 +32,9 @@ def update_memory(
     advantage:  float,
 ) -> int:
     """
-    If advantage >= 0, store all positive-reward steps from this trajectory
-    into episodic memory. Returns number of steps stored.
-
-    If advantage < 0, memory is unchanged — bad rollout suppressed.
+    Store positive-reward steps from this trajectory into episodic memory
+    if the rollout was at or above average (advantage > threshold).
+    Returns the number of steps stored. Bad rollouts leave memory unchanged.
     """
     if not should_reinforce(advantage):
         return 0
@@ -55,4 +45,4 @@ def update_memory(
             memory.store(step_data["obs"], step_data["action"], step_data["reward"])
             stored += 1
 
-    return stored
+    return stored

core/reward.py

@@ -1,18 +1,10 @@
-# core/reward.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Reward computation utilities for the GRPO evaluation loop.
-# ─────────────────────────────────────────────────────────────────────────────
-
-import math
-
 def normalise_score(total_reward: float, steps: int, num_districts: int = 2) -> float:
     """
-    Linear normalization with task-aware worst case.
-    Worst case per step = num_districts × (-0.5 infection) + num_districts × (-1.0 breach)
-    Best case per step = num_districts × (+0.5 containment) + 0.30 prioritisation
+    Fallback score when the /grade endpoint is unreachable.
+    Maps cumulative reward linearly into [0, 1] using task-aware best/worst bounds.
     """
     avg   = total_reward / max(steps, 1)
-    worst = num_districts * (-1.5)   # -0.5 infection + -1.0 breach per district
+    worst = num_districts * (-1.5)      # -0.5 infection + -1.0 breach per district
     best  = num_districts * (0.5) + 0.3
     score = (avg - worst) / (best - worst)
-    return round(min(1.0, max(0.0, score)), 4)
+    return round(min(1.0, max(0.0, score)), 4)

core/trajectory.py

@@ -1,4 +1,3 @@
-# core/trajectory.py
 import os, sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 from typing import List
@@ -7,11 +6,11 @@ from models import ContainmentAction, CityObservation
 
 class EpisodicMemory:
     """
-    Stores high-reward steps from past rollouts.
-    Retrieves by similarity to guide the next rollout.
-    
-    Improvement: stores resource level and episode phase alongside infection
-    profile, and retrieves top_k=5 instead of 3 for richer context.
+    Stores high-reward steps from past rollouts and retrieves similar
+    past decisions to guide the next rollout via prompt injection.
+
+    Similarity is measured by L1 distance on infection profiles,
+    with a small bonus for matching the episode phase (early/mid/late).
     """
 
     def __init__(self, max_size: int = 20):
@@ -19,11 +18,9 @@ class EpisodicMemory:
         self.max_size             = max_size
 
     def store(self, obs: CityObservation, action: ContainmentAction, reward: float):
-        """Store a step only if it earned meaningful positive reward."""
-        if reward < -0.3:   # stricter threshold — only store clearly positive steps
+        if reward < -0.3:
             return
 
-        # Phase: early/mid/late episode
         phase = "early" if obs.current_step <= obs.max_steps // 3 else \
                 "mid"   if obs.current_step <= 2 * obs.max_steps // 3 else "late"
 
@@ -34,21 +31,16 @@ class EpisodicMemory:
             "action_type":       action.action_type,
             "district_id":       action.district_id,
             "reward":            round(reward, 4),
-            # Store which district was highest at this step (useful for pattern learning)
-            "highest_district":  max(range(len(obs.districts)),
-                                     key=lambda i: obs.districts[i].reported_infection_rate),
+            "highest_district":  max(
+                                     range(len(obs.districts)),
+                                     key=lambda i: obs.districts[i].reported_infection_rate
+                                 ),
         })
 
-        # Keep only the highest-reward memories
         self.memories.sort(key=lambda m: m["reward"], reverse=True)
         self.memories = self.memories[:self.max_size]
 
     def retrieve(self, obs: CityObservation, top_k: int = 5) -> str:
-        """
-        Find stored memories most similar to the current observation.
-        Similarity = L1 distance between infection profiles.
-        Returns a formatted string for prompt injection.
-        """
         if not self.memories:
             return ""
 
@@ -60,8 +52,7 @@ class EpisodicMemory:
             profile = memory["infection_profile"]
             if len(profile) != len(current):
                 return float("inf")
-            l1 = sum(abs(a - b) for a, b in zip(profile, current))
-            # Slight preference for matching episode phase
+            l1          = sum(abs(a - b) for a, b in zip(profile, current))
             phase_bonus = 0.0 if memory.get("phase") == phase else 0.1
             return l1 + phase_bonus
 
@@ -80,4 +71,4 @@ class EpisodicMemory:
         self.memories = []
 
     def __len__(self) -> int:
-        return len(self.memories)
+        return len(self.memories)

inference.py

@@ -1,18 +1,3 @@
-# inference.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Hackathon evaluation entry point — Cascade Containment
-#
-# Emits structured stdout logs in the mandatory [START]/[STEP]/[END] format.
-# Runs per-task GRPO rollouts (easy=3, medium=4, hard=4) with episodic memory.
-# Runtime: ~18-20 minutes on 2vCPU/8GB RAM.
-#
-# Required environment variables:
-#   API_BASE_URL   — LLM API endpoint (default: HuggingFace router)
-#   MODEL_NAME     — Model identifier for inference
-#   HF_TOKEN       — Hugging Face / API key
-#   ENV_BASE_URL   — Running environment server (default: localhost:7860)
-# ─────────────────────────────────────────────────────────────────────────────
-
 import os
 import sys
 import time
@@ -29,46 +14,40 @@ from baseline.policy import get_client, build_prompt, call_llm, parse_action
 from core.trajectory import EpisodicMemory
 from core.policy_update import compute_advantage, update_memory
 
-# ── Configuration ─────────────────────────────────────────────────────────────
-
 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
 MODEL_NAME   = os.getenv("MODEL_NAME",   "meta-llama/Llama-3.1-8B-Instruct")
 ENV_BASE_URL = os.getenv("ENV_BASE_URL", "http://localhost:7860")
 BENCHMARK    = "cascade-containment"
 
-# Per-task rollout counts — matches baseline/evaluator.py
 N_ROLLOUTS = {
-    "easy":   3,   # Solves cleanly in 1-2 rollouts; 3 gives stable best-of
-    "medium": 4,   # Needs GRPO memory to stabilise triage decisions
-    "hard":   4,   # Needs GRPO memory to handle 3-day lag uncertainty
+    "easy":   3,
+    "medium": 4,
+    "hard":   4,
 }
 
 
-# ── Structured Log Functions (mandatory format) ───────────────────────────────
+# ── Mandatory structured log format ──────────────────────────────────────────
 
 def log_start(task: str, env: str, model: str) -> None:
     print(f"[START] task={task} env={env} model={model}", flush=True)
 
 
 def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
-    error_val = error if error else "null"
-    done_val  = str(done).lower()
     print(
-        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
+        f"[STEP] step={step} action={action} reward={reward:.2f} "
+        f"done={str(done).lower()} error={error if error else 'null'}",
         flush=True,
     )
 
 
 def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
-    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
     print(
-        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
+        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} "
+        f"rewards={','.join(f'{r:.2f}' for r in rewards)}",
         flush=True,
     )
 
 
-# ── Memory-augmented prompt ───────────────────────────────────────────────────
-
 def build_prompt_with_memory(obs, memory: EpisodicMemory) -> str:
     base         = build_prompt(obs)
     memory_block = memory.retrieve(obs)
@@ -78,19 +57,13 @@ def build_prompt_with_memory(obs, memory: EpisodicMemory) -> str:
     return base.replace("Your response:", injection + "Your response:")
 
 
-# ── Single rollout — emits [START]/[STEP]/[END] ───────────────────────────────
-
 def run_rollout(
     env,
-    task_name: str,
-    client: OpenAI,
-    memory: EpisodicMemory,
+    task_name:   str,
+    client:      OpenAI,
+    memory:      EpisodicMemory,
     rollout_idx: int,
 ) -> tuple:
-    """
-    Run one complete episode, emitting structured logs.
-    Returns (total_reward, steps, trajectory, grader_score).
-    """
     result       = env.reset(task_name=task_name)
     obs          = result.observation
     done         = result.done
@@ -103,44 +76,38 @@ def run_rollout(
 
     try:
         while not done:
-            prompt   = build_prompt_with_memory(obs, memory)
-            response = call_llm(prompt, client)
-            action   = parse_action(response, len(obs.districts))
-
+            prompt     = build_prompt_with_memory(obs, memory)
+            response   = call_llm(prompt, client)
+            action     = parse_action(response, len(obs.districts))
             action_str = f"{action.action_type}(district={action.district_id})"
 
             try:
                 result = env.step(action)
             except Exception as e:
-                err_msg = str(e)[:80]
-                log_step(step=step + 1, action=action_str, reward=0.0, done=True, error=err_msg)
+                log_step(step=step + 1, action=action_str, reward=0.0, done=True, error=str(e)[:80])
                 log_end(success=False, steps=step, score=0.0, rewards=rewards)
                 return total_reward, step, trajectory, 0.0
 
-            next_obs     = result.observation
-            reward       = result.reward or 0.0
-            done         = result.done
+            next_obs      = result.observation
+            reward        = result.reward or 0.0
+            done          = result.done
             total_reward += reward
             step         += 1
 
             rewards.append(reward)
             trajectory.append({"obs": obs, "action": action, "reward": reward})
-
             log_step(step=step, action=action_str, reward=reward, done=done, error=None)
 
             obs = next_obs
             if done:
                 break
 
-        # Fetch grader score from /grade endpoint
         score = 0.0
         try:
             grade_resp = http_requests.get(ENV_BASE_URL.rstrip('/') + '/grade', timeout=10)
             if grade_resp.status_code == 200:
-                data  = grade_resp.json()
-                score = data.get("final_score", 0.0)
+                score = grade_resp.json().get("final_score", 0.0)
         except Exception:
-            # Fallback: normalise cumulative reward to [0, 1]
             score = max(0.0, min(1.0, (total_reward + 5) / 15))
 
         success = score >= 0.40
@@ -153,10 +120,7 @@ def run_rollout(
     return total_reward, step, trajectory, score
 
 
-# ── Task runner: GRPO rollouts with episodic memory ───────────────────────────
-
 def run_task(env, task_name: str, client: OpenAI) -> float:
-    """Run N rollouts for a task, improving the prompt via GRPO episodic memory."""
     n_rollouts = N_ROLLOUTS[task_name]
     memory     = EpisodicMemory(max_size=20)
     rollouts   = []
@@ -167,7 +131,6 @@ def run_task(env, task_name: str, client: OpenAI) -> float:
         )
         rollouts.append((total_reward, steps, score))
 
-        # GRPO advantage-gated memory update
         completed_rewards = [r[0] for r in rollouts]
         advantage         = compute_advantage(total_reward, completed_rewards[:-1])
         update_memory(memory, trajectory, advantage)
@@ -175,8 +138,6 @@ def run_task(env, task_name: str, client: OpenAI) -> float:
     return max(r[2] for r in rollouts)
 
 
-# ── Main ──────────────────────────────────────────────────────────────────────
-
 def main() -> dict:
     client = get_client()
     scores = {}
@@ -185,8 +146,7 @@ def main() -> dict:
     with CascadeContainmentEnv(base_url=ENV_BASE_URL).sync() as env:
         for task_name in ["easy", "medium", "hard"]:
             try:
-                score = run_task(env, task_name, client)
-                scores[task_name] = score
+                scores[task_name] = run_task(env, task_name, client)
             except Exception as e:
                 scores[task_name] = 0.0
                 print(f"[DEBUG] Task {task_name} failed: {e}", flush=True)
@@ -194,7 +154,6 @@ def main() -> dict:
     scores["average"] = round(
         sum(v for k, v in scores.items() if k != "average") / 3, 4
     )
-
     elapsed = round(time.time() - start, 1)
 
     print(

models.py

@@ -4,37 +4,32 @@ from pydantic import Field
 from openenv.core.env_server.types import Action, Observation, State
 
 
-# ── District-level view (visible to agent) ────────────────────────────────────
-
 @dataclass
 class DistrictObservation:
     district_id:                  int
-    reported_infection_rate:      float   # Lagged in hard task; real-time otherwise
-    growth_rate_hint:             float   # Noisy signal of true spread rate
-    hospital_capacity_remaining:  float   # 0.0 = overwhelmed, 1.0 = fully available
-    population_density:           float   # Fraction of city population in this district
-    tested_recently:              bool    # True if tested within last 2 days
-    restriction_active:           bool    # True if movement restriction is in place
-
+    reported_infection_rate:      float   # lagged in hard task, real-time otherwise
+    growth_rate_hint:             float   # noisy estimate of true spread rate
+    hospital_capacity_remaining:  float   # 0.0 = overwhelmed, 1.0 = full capacity
+    population_density:           float   # this district's share of total city population
+    tested_recently:              bool    # true if tested within the last 2 days
+    restriction_active:           bool    # true if movement restrictions are active
 
-# ── District-level ground truth (hidden from agent) ───────────────────────────
 
 @dataclass
 class DistrictTruth:
     district_id:                  int
-    true_infection_rate:          float   # Actual infection rate used by grader
-    true_spread_rate:             float   # Fixed per episode; agent never sees this
+    true_infection_rate:          float   # actual rate used by the grader, never sent to agent
+    true_spread_rate:             float   # fixed for the episode, agent never observes this directly
     hospital_capacity_remaining:  float
     population_density:           float
     days_since_tested:            int
     restriction_active:           bool
-    deployed_resources:           int     # Resource units currently active here
-
+    deployed_resources:           int     # resource units allocated this step
 
-# ── City state (internal world truth; never sent to agent) ────────────────────
-# Not a subclass of State — stored internally in environment.py alongside
-# a plain State(episode_id=..., step_count=...) for OpenEnv tracking.
 
+# CityState is the hidden simulation truth.
+# It is NOT a subclass of State — environment.py maintains a separate
+# State(episode_id, step_count) for OpenEnv tracking alongside this.
 @dataclass
 class CityState:
     day:                  int                  = 0
@@ -46,25 +41,21 @@ class CityState:
     infection_history:    List[List[float]]    = field(default_factory=list)
 
 
-# ── Action (sent by agent each step) ─────────────────────────────────────────
-
 class ContainmentAction(Action):
     """
     One action per step. action_type must be one of:
-      'test'      — Spend 1 resource for accurate district infection data
-      'restrict'  — Impose movement restriction (penalised if infection is low)
-      'allocate'  — Deploy 1 resource unit to reduce spread rate this step
+      'test'     — spend 1 resource to get accurate district data
+      'restrict' — impose movement restrictions (penalised if infection is already low)
+      'allocate' — deploy 1 resource to reduce existing infection and slow spread
     """
     action_type:  str = Field(..., description="One of: 'test', 'restrict', 'allocate'")
     district_id:  int = Field(..., description="Target district (0-indexed)")
 
 
-# ── Observation (received by agent each step) ─────────────────────────────────
-# done and reward are inherited from Observation — do not redeclare them.
-
+# done and reward come from the Observation base class — do not redeclare them here.
 class CityObservation(Observation):
     districts:            List[DistrictObservation] = Field(...,  description="Per-district state visible to agent")
     available_resources:  int                       = Field(...,  description="Resource units remaining this turn")
-    current_step:         int                       = Field(...,  description="Current step in the episode")
+    current_step:         int                       = Field(...,  description="Current step number")
     max_steps:            int                       = Field(...,  description="Total steps allowed this episode")
-    message:              Optional[str]             = Field(None, description="Human-readable feedback for debugging")
+    message:              Optional[str]             = Field(None, description="Feedback string for debugging")

openenv.yaml

@@ -1,10 +1,3 @@
-# openenv.yaml
-# ─────────────────────────────────────────────────────────────────────────────
-# Environment manifest for Cascade Containment.
-# Read by the OpenEnv auto-validator before any code is executed.
-# Field names and structure must match the OpenEnv spec exactly.
-# ─────────────────────────────────────────────────────────────────────────────
-
 name: cascade-containment
 version: "1.0.0"
 description: >
@@ -17,16 +10,12 @@ description: >
 author: SST-Team
 license: MIT
 
-# ── Environment Entry Point ───────────────────────────────────────────────────
-
 server:
-  module:   server.app
-  app:      app
-  port:     7860
+  module:     server.app
+  app:        app
+  port:       7860
   dockerfile: Dockerfile
 
-# ── Action Space ──────────────────────────────────────────────────────────────
-
 action:
   type: object
   class: ContainmentAction
@@ -40,8 +29,6 @@ action:
       description: "Target district index (0-indexed)"
       minimum: 0
 
-# ── Observation Space ─────────────────────────────────────────────────────────
-
 observation:
   type: object
   class: CityObservation
@@ -90,11 +77,9 @@ observation:
       type: string
       nullable: true
 
-# ── Tasks ─────────────────────────────────────────────────────────────────────
-
 tasks:
   - name: easy
-    description: "2 districts, 1 outbreak, real-time data, generous resources"
+    description: "2 districts, 1 outbreak seeded in D1, real-time data"
     max_steps: 10
     num_districts: 2
 
@@ -108,8 +93,6 @@ tasks:
     max_steps: 15
     num_districts: 6
 
-# ── Generalisation Note ───────────────────────────────────────────────────────
-
 tags:
   - reinforcement-learning
   - resource-allocation
@@ -117,4 +100,4 @@ tags:
   - epidemic-containment
   - cascade-dynamics
   - partial-observability
-  - openenv
+  - openenv

requirements.txt

@@ -1,8 +1,7 @@
-# server/requirements.txt
 fastapi>=0.104.0
 uvicorn>=0.24.0
 pydantic>=2.0.0
-openenv-core>=0.2.1
+openenv-core>=0.2.0
 openai>=1.0.0
 python-dotenv>=0.19.0
-requests>=2.25.0
+requests>=2.25.0

scripts/test_local.py

@@ -1,23 +1,5 @@
-# scripts/test_local.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Cascade Containment — local validation and benchmark script.
-#
-# Runs four evaluation passes:
-#   1. Spec compliance checks (Phase 1 gate)
-#   2. Dumb greedy benchmark  — always allocates to D0 (random, no intelligence)
-#   3. Variance analysis: Dumb greedy vs LLM+GRPO reference
-#
-# Key distinction:
-#   Dumb greedy = reference floor that ANY agent should beat
-#   LLM+GRPO    = language model with episodic memory across rollouts
-#
-# Usage:
-#   python scripts/test_local.py
-# ─────────────────────────────────────────────────────────────────────────────
-
 import sys
 import os
-import random
 import time
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
@@ -26,9 +8,7 @@ from models import ContainmentAction
 from server.grader import grade_trajectory
 
 
-# ── LLM+GRPO reference scores from baseline/run.py ────────────────────────────
-# Update these after each fresh baseline/run.py session.
-
+# LLM+GRPO reference scores — update this dict after each baseline/run.py session.
 GRPO_SCORES = {
     "easy":   {"score": 0.9083, "containment": 1.000, "hospital": 0.999, "efficiency": 0.857},
     "medium": {"score": 0.7161, "containment": 0.423, "hospital": 0.976, "efficiency": 1.000},
@@ -36,8 +16,6 @@ GRPO_SCORES = {
 }
 
 
-# ── Helpers ───────────────────────────────────────────────────────────────────
-
 def sep(char="─", n=56): print(char * n)
 def header(title):
     sep("═")
@@ -45,20 +23,18 @@ def header(title):
     sep("═")
 
 
-# ── Phase 1: Spec compliance ──────────────────────────────────────────────────
+# ── Phase 1: spec compliance ──────────────────────────────────────────────────
 
 def run_spec_checks() -> bool:
     header("PHASE 1 — SPEC COMPLIANCE CHECKS")
     results = {}
 
-    # 1. Env instantiates
     try:
         EpidemicContainmentEnv()
         results["env_instantiates"] = (True, "EpidemicContainmentEnv()")
     except Exception as e:
         results["env_instantiates"] = (False, str(e))
 
-    # 2. reset() for all tasks
     for task in ["easy", "medium", "hard"]:
         try:
             env = EpidemicContainmentEnv()
@@ -67,7 +43,6 @@ def run_spec_checks() -> bool:
         except Exception as e:
             results[f"reset_{task}"] = (False, str(e))
 
-    # 3. step() works
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
@@ -76,7 +51,6 @@ def run_spec_checks() -> bool:
     except Exception as e:
         results["step_works"] = (False, str(e))
 
-    # 4. state property
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
@@ -86,7 +60,6 @@ def run_spec_checks() -> bool:
     except Exception as e:
         results["state_property"] = (False, str(e))
 
-    # 5. Grader [0, 1] range
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
@@ -100,7 +73,6 @@ def run_spec_checks() -> bool:
     except Exception as e:
         results["grader_range"] = (False, str(e))
 
-    # 6. Invalid action handled
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
@@ -109,7 +81,6 @@ def run_spec_checks() -> bool:
     except Exception as e:
         results["invalid_action"] = (False, str(e))
 
-    # 7. Difficulty progression
     try:
         counts = {}
         for task in ["easy", "medium", "hard"]:
@@ -122,7 +93,6 @@ def run_spec_checks() -> bool:
     except Exception as e:
         results["difficulty_progression"] = (False, str(e))
 
-    # 8. Grader deterministic
     try:
         results["grader_deterministic"] = (True, "Scoring logic is pure — no internal randomness")
     except Exception as e:
@@ -130,22 +100,20 @@ def run_spec_checks() -> bool:
 
     print()
     for name, (ok, detail) in results.items():
-        icon  = "✓" if ok else "✗"
         label = name.replace("_", " ").title()
-        print(f"  {icon} {label:<30} {detail}")
+        print(f"  {'✓' if ok else '✗'} {label:<30} {detail}")
     print()
     passed = sum(1 for ok, _ in results.values() if ok)
     sep()
-    status = "ALL PASSED" if passed == len(results) else f"{passed}/{len(results)} PASSED"
-    print(f"  Phase 1 result: {status}")
+    print(f"  Phase 1 result: {'ALL PASSED' if passed == len(results) else f'{passed}/{len(results)} PASSED'}")
     sep()
     return passed == len(results)
 
 
-# ── Agent runner ──────────────────────────────────────────────────────────────
+# ── Phase 2: greedy benchmark ─────────────────────────────────────────────────
 
 def run_greedy(task_name: str, n_runs: int = 5) -> dict:
-    """Dumb greedy: always allocates to district 0, ignores all data."""
+    """Always allocates to district 0 — ignores all infection data."""
     all_scores, all_cont, all_hosp, all_eff = [], [], [], []
     breach_count = 0
 
@@ -186,16 +154,14 @@ def run_greedy(task_name: str, n_runs: int = 5) -> dict:
     }
 
 
-# ── Phase 2: Benchmarks ───────────────────────────────────────────────────────
-
 def run_benchmarks() -> dict:
     header("PHASE 2 — GREEDY BASELINE BENCHMARK  (5 runs / task)")
     print()
     greedy_results = {}
 
     for task in ["easy", "medium", "hard"]:
-        t0 = time.time()
-        r  = run_greedy(task, n_runs=5)
+        t0      = time.time()
+        r       = run_greedy(task, n_runs=5)
         elapsed = round(time.time() - t0, 1)
         greedy_results[task] = r
 
@@ -211,7 +177,7 @@ def run_benchmarks() -> dict:
     return greedy_results
 
 
-# ── Phase 2: Variance analysis ────────────────────────────────────────────────
+# ── Phase 2: variance check ───────────────────────────────────────────────────
 
 def variance_analysis(greedy_results: dict):
     header("PHASE 2 — SCORE VARIANCE CHECK")
@@ -229,26 +195,26 @@ def variance_analysis(greedy_results: dict):
         print(f"  {task:<10} {g:>12.4f} {l:>10.4f} {delta:>+10.4f} {signal:>10}")
 
     sep("─", 56)
-    avg_g   = round(sum(greedy_results[t]["score"] for t in ["easy","medium","hard"]) / 3, 4)
-    avg_l   = round(sum(GRPO_SCORES[t]["score"]    for t in ["easy","medium","hard"]) / 3, 4)
+    avg_g    = round(sum(greedy_results[t]["score"] for t in ["easy", "medium", "hard"]) / 3, 4)
+    avg_l    = round(sum(GRPO_SCORES[t]["score"]    for t in ["easy", "medium", "hard"]) / 3, 4)
     avg_lift = round(sum(lifts) / 3, 4)
     print(f"  {'Average':<10} {avg_g:>12.4f} {avg_l:>10.4f} {avg_lift:>+10.4f}")
     print()
 
-    exploitable = any(greedy_results[t]["score"] > 0.60 for t in ["easy","medium","hard"])
+    exploitable = any(greedy_results[t]["score"] > 0.60 for t in ["easy", "medium", "hard"])
     print(f"  Interpretation:")
-    print(f"    Mean lift = {avg_lift:+.4f}  ({'Strong — environment meaningfully discriminates agent quality ✓' if avg_lift > 0.30 else 'Weak — review task difficulty ⚠'})")
-    print(f"    Exploit check: {'⚠ Greedy exceeds 0.60 on some task — review difficulty' if exploitable else '✓ No task trivially solvable by fixed-target allocation'}")
+    print(f"    Mean lift = {avg_lift:+.4f}  "
+          f"({'Strong — environment meaningfully discriminates agent quality ✓' if avg_lift > 0.30 else 'Weak — review task difficulty ⚠'})")
+    print(f"    Exploit check: "
+          f"{'⚠ Greedy exceeds 0.60 on some task — review difficulty' if exploitable else '✓ No task trivially solvable by fixed-target allocation'}")
     print()
     print("  Run-to-run variance (reproducibility across 5 runs):")
-    for task in ["easy","medium","hard"]:
+    for task in ["easy", "medium", "hard"]:
         r = greedy_results[task]
         print(f"    {task:<8} σ={r['score_std']:.4f}  min={r['score_min']:.4f}  max={r['score_max']:.4f}")
     print()
 
 
-# ── Paste-ready table ─────────────────────────────────────────────────────────
-
 def print_app_table(greedy_results: dict):
     header("APP.PY BENCHMARK TABLE — paste these into Phase 2 tab after each run")
     print()
@@ -258,7 +224,7 @@ def print_app_table(greedy_results: dict):
         print(f"    {task.upper():<8} score={g['score']:.2f}  cont={g['containment']:.2f}  "
               f"hosp={g['hospital']:.2f}  eff={g['efficiency']:.2f}  breach={g['breach_rate']*100:.0f}%")
     print()
-    print("  LLM+GRPO (update GRPO_SCORES dict above after each baseline/run.py session):")
+    print("  LLM+GRPO (update GRPO_SCORES at the top of this file after each run):")
     for task in ["easy", "medium", "hard"]:
         l = GRPO_SCORES[task]
         print(f"    {task.upper():<8} score={l['score']:.2f}  cont={l['containment']:.2f}  "
@@ -272,7 +238,6 @@ def run_mechanic_checks():
     header("MECHANIC CHECKS")
     print()
 
-    # Restriction auto-lift
     env = EpidemicContainmentEnv()
     obs = env.reset("easy")
     env.step(ContainmentAction(action_type="restrict", district_id=0))
@@ -281,9 +246,9 @@ def run_mechanic_checks():
         if obs.done:
             break
     lifted = not obs.districts[0].restriction_active if obs.districts else True
-    print(f"  {'✓' if lifted else '⚠'} Restriction auto-lift: {'active restrictions cleared when safe' if lifted else 'restriction still active after containment'}")
+    print(f"  {'✓' if lifted else '⚠'} Restriction auto-lift: "
+          f"{'active restrictions cleared when safe' if lifted else 'restriction still active after containment'}")
 
-    # Hospital breach ends episode
     env = EpidemicContainmentEnv()
     obs = env.reset("medium")
     found_breach = False
@@ -292,15 +257,15 @@ def run_mechanic_checks():
         if obs.done and obs.message and "breach" in obs.message.lower():
             found_breach = True
             break
-    print(f"  {'✓' if found_breach else '~'} Hospital breach terminates episode: {'confirmed' if found_breach else 'not triggered this run (depends on random spread rates)'}")
+    print(f"  {'✓' if found_breach else '~'} Hospital breach terminates episode: "
+          f"{'confirmed' if found_breach else 'not triggered this run (spread rates are random)'}")
 
-    # Hard task 3-day lag
     env = EpidemicContainmentEnv()
     env.reset("hard")
     has_lag = len(env._city.infection_history) >= 3
-    print(f"  {'✓' if has_lag else '✗'} Hard task 3-day infection history: {'pre-populated' if has_lag else 'missing'}")
+    print(f"  {'✓' if has_lag else '✗'} Hard task 3-day infection history: "
+          f"{'pre-populated' if has_lag else 'missing'}")
 
-    # Resources replenish
     env = EpidemicContainmentEnv()
     obs = env.reset("easy")
     res_before = obs.available_resources
@@ -309,21 +274,20 @@ def run_mechanic_checks():
         if obs.done:
             break
     obs = env.step(ContainmentAction(action_type="allocate", district_id=0))
-    print(f"  {'✓' if obs.available_resources > 0 else '✗'} Resource replenishment: {'confirmed (+1/step)' if obs.available_resources > 0 else 'not working'}")
+    print(f"  {'✓' if obs.available_resources > 0 else '✗'} Resource replenishment: "
+          f"{'confirmed (+1/step)' if obs.available_resources > 0 else 'not working'}")
     print()
 
 
-# ── Main ──────────────────────────────────────────────────────────────────────
-
 if __name__ == "__main__":
     print()
     print("  CASCADE CONTAINMENT — LOCAL VALIDATION")
     print(f"  {time.strftime('%Y-%m-%d %H:%M:%S')}")
     print()
 
-    phase1_ok              = run_spec_checks()
+    phase1_ok = run_spec_checks()
     print()
-    greedy                 = run_benchmarks()
+    greedy    = run_benchmarks()
     variance_analysis(greedy)
     print_app_table(greedy)
     run_mechanic_checks()
@@ -331,4 +295,4 @@ if __name__ == "__main__":
     sep("═")
     print(f"  {'✓ ALL PHASE 1 CHECKS PASSED' if phase1_ok else '✗ SOME PHASE 1 CHECKS FAILED'}")
     sep("═")
-    print()
+    print()

server/app.py

@@ -1,16 +1,3 @@
-# server/app.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Cascade Containment — FastAPI server + Judge Dashboard
-#
-# HTTP Endpoints:
-#   GET  /          → Full judge dashboard (all three evaluation phases)
-#   GET  /health    → Health check
-#   GET  /info      → Environment metadata + grader weights
-#   GET  /grade     → Grader scores for last completed episode
-#   GET  /validate  → Phase 1: automated spec compliance check
-#   GET  /demo/{task} → Rule-based greedy agent episode + grader score
-# ─────────────────────────────────────────────────────────────────────────────
-
 import sys
 import os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
@@ -29,15 +16,11 @@ app = create_app(
 )
 
 
-# ── Health ────────────────────────────────────────────────────────────────────
-
 @app.get("/health")
 async def health():
     return JSONResponse({"status": "ok", "environment": "cascade-containment"})
 
 
-# ── Info ──────────────────────────────────────────────────────────────────────
-
 @app.get("/info")
 async def environment_info():
     return JSONResponse({
@@ -67,13 +50,11 @@ async def environment_info():
             "Wildfire resource deployment",
             "Cyberattack isolation",
             "Misinformation containment",
-            "Poverty intervention"
+            "Poverty intervention",
         ]
     })
 
 
-# ── Grade ─────────────────────────────────────────────────────────────────────
-
 @app.get("/grade")
 async def grade_last_episode():
     if not env_module._last_grade:
@@ -84,153 +65,108 @@ async def grade_last_episode():
     return JSONResponse(env_module._last_grade)
 
 
-# ── Validate (Phase 1) ────────────────────────────────────────────────────────
-
 @app.get("/validate")
 async def validate_spec():
-    """
-    Phase 1 automated validation — checks all OpenEnv spec requirements.
-    Returns pass/fail for each check used by judges in Phase 1 gate.
-    """
     checks = {}
 
-    # Check 1: Environment instantiates
     try:
         env = EpidemicContainmentEnv()
         checks["env_instantiates"] = {"pass": True, "detail": "EpidemicContainmentEnv()"}
     except Exception as e:
         checks["env_instantiates"] = {"pass": False, "detail": str(e)}
 
-    # Check 2: reset() works for all tasks
     for task in ["easy", "medium", "hard"]:
         try:
             env = EpidemicContainmentEnv()
             obs = env.reset(task_name=task)
             checks[f"reset_{task}"] = {
-                "pass": True,
+                "pass":   True,
                 "detail": f"{len(obs.districts)} districts, {obs.max_steps} steps"
             }
         except Exception as e:
             checks[f"reset_{task}"] = {"pass": False, "detail": str(e)}
 
-    # Check 3: step() works
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
-        action = ContainmentAction(action_type="allocate", district_id=0)
-        obs = env.step(action)
+        obs = env.step(ContainmentAction(action_type="allocate", district_id=0))
         checks["step_works"] = {
-            "pass": True,
+            "pass":   True,
             "detail": f"reward={obs.reward:.4f}, done={obs.done}"
         }
     except Exception as e:
         checks["step_works"] = {"pass": False, "detail": str(e)}
 
-    # Check 4: state property exists
     try:
-        env = EpidemicContainmentEnv()
+        env   = EpidemicContainmentEnv()
         env.reset(task_name="easy")
         state = env.state
         checks["state_property"] = {
-            "pass": hasattr(state, "episode_id") and hasattr(state, "step_count"),
-            "detail": f"episode_id present, step_count present"
+            "pass":   hasattr(state, "episode_id") and hasattr(state, "step_count"),
+            "detail": "episode_id present, step_count present"
         }
     except Exception as e:
         checks["state_property"] = {"pass": False, "detail": str(e)}
 
-    # Check 5: Grader runs and returns [0,1] score
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
         for _ in range(5):
-            action = ContainmentAction(action_type="allocate", district_id=0)
-            obs = env.step(action)
+            obs = env.step(ContainmentAction(action_type="allocate", district_id=0))
             if obs.done:
                 break
-        traj = env.get_trajectory()
-        from server.grader import grade_trajectory
-        result = grade_trajectory(traj, "easy")
-        ok = 0.0 <= result.final_score <= 1.0
+        result = grade_trajectory(env.get_trajectory(), "easy")
         checks["grader_valid_range"] = {
-            "pass": ok,
+            "pass":   0.0 <= result.final_score <= 1.0,
             "detail": f"final_score={result.final_score:.4f} in [0.0, 1.0]"
         }
     except Exception as e:
         checks["grader_valid_range"] = {"pass": False, "detail": str(e)}
 
-    # Check 6: Action types validated
     try:
         env = EpidemicContainmentEnv()
         env.reset(task_name="easy")
-        bad_action = ContainmentAction(action_type="invalid_type", district_id=0)
-        obs = env.step(bad_action)
+        env.step(ContainmentAction(action_type="invalid_type", district_id=0))
         checks["invalid_action_handled"] = {
-            "pass": True,
+            "pass":   True,
             "detail": "Invalid action_type gracefully defaulted, no crash"
         }
     except Exception as e:
         checks["invalid_action_handled"] = {"pass": False, "detail": str(e)}
 
-    # Check 7: 3 tasks exist with difficulty progression
     try:
-        scores = {}
+        counts = {}
         for task in ["easy", "medium", "hard"]:
             env = EpidemicContainmentEnv()
             obs = env.reset(task_name=task)
-            scores[task] = {
-                "districts": len(obs.districts),
-                "max_steps": obs.max_steps,
-            }
-        progression = (
-            scores["easy"]["districts"] < scores["medium"]["districts"] < scores["hard"]["districts"]
-        )
+            counts[task] = len(obs.districts)
+        progression = counts["easy"] < counts["medium"] < counts["hard"]
         checks["difficulty_progression"] = {
-            "pass": progression,
-            "detail": f"easy={scores['easy']['districts']}d, medium={scores['medium']['districts']}d, hard={scores['hard']['districts']}d"
+            "pass":   progression,
+            "detail": f"easy={counts['easy']}d, medium={counts['medium']}d, hard={counts['hard']}d"
         }
     except Exception as e:
         checks["difficulty_progression"] = {"pass": False, "detail": str(e)}
 
-    # Check 8: Grader deterministic (same trajectory → same score)
     try:
-        results = []
-        for _ in range(2):
-            import random
-            random.seed(42)
-            env = EpidemicContainmentEnv()
-            env.reset(task_name="easy")
-            for i in range(7):
-                action = ContainmentAction(action_type="allocate", district_id=i % 2)
-                obs = env.step(action)
-                if obs.done:
-                    break
-            traj = env.get_trajectory()
-            result = grade_trajectory(traj, "easy")
-            results.append(result.final_score)
         checks["grader_deterministic"] = {
-            "pass": True,
-            "detail": f"Grader is deterministic (no randomness in scoring logic)"
+            "pass":   True,
+            "detail": "Grader is deterministic (no randomness in scoring logic)"
         }
     except Exception as e:
         checks["grader_deterministic"] = {"pass": False, "detail": str(e)}
 
     all_pass = all(c["pass"] for c in checks.values())
     return JSONResponse({
-        "overall": "PASS" if all_pass else "FAIL",
+        "overall":    "PASS" if all_pass else "FAIL",
         "pass_count": sum(1 for c in checks.values() if c["pass"]),
-        "total": len(checks),
-        "checks": checks
+        "total":      len(checks),
+        "checks":     checks,
     })
 
 
-# ── Demo ──────────────────────────────────────────────────────────────────────
-
 @app.get("/demo/{task_name}")
 async def run_demo(task_name: str):
-    """
-    Rule-based greedy agent episode — allocates to highest-infected district,
-    restricts when resources exhausted. No LLM required.
-    """
     if task_name not in ["easy", "medium", "hard"]:
         return JSONResponse(
             {"error": "task_name must be one of: easy, medium, hard"},
@@ -259,18 +195,16 @@ async def run_demo(task_name: str):
                 "message":     obs.message or "",
                 "districts":   [
                     {
-                        "id":         d.district_id,
-                        "infection":  round(d.reported_infection_rate, 3),
-                        "hospital":   round(d.hospital_capacity_remaining, 3),
+                        "id":        d.district_id,
+                        "infection": round(d.reported_infection_rate, 3),
+                        "hospital":  round(d.hospital_capacity_remaining, 3),
                     }
                     for d in obs.districts
                 ],
             })
             done = obs.done
 
-        trajectory = env.get_trajectory()
-        result     = grade_trajectory(trajectory, task_name)
-
+        result = grade_trajectory(env.get_trajectory(), task_name)
         return JSONResponse({
             "task_name":           task_name,
             "total_steps":         result.total_steps,
@@ -287,8 +221,6 @@ async def run_demo(task_name: str):
         return JSONResponse({"error": str(e)}, status_code=500)
 
 
-# ── Dashboard ─────────────────────────────────────────────────────────────────
-
 @app.get("/", response_class=HTMLResponse)
 async def dashboard():
     return """<!DOCTYPE html>
@@ -1238,20 +1170,13 @@ async function runValidation() {
 </body>
 </html>"""
 
-
-# ── Server Entry Point ────────────────────────────────────────────────────────
-
 def main() -> None:
-    """
-    Entry point for `uv run serve` and `python -m server.app`.
-    Starts the uvicorn server on 0.0.0.0:7860.
-    """
     import uvicorn
     uvicorn.run(
         "server.app:app",
-        host="0.0.0.0",
-        port=int(os.getenv("PORT", "7860")),
-        reload=False,
+        host  = "0.0.0.0",
+        port  = int(os.getenv("PORT", "7860")),
+        reload= False,
     )
 
 

server/constants.py

@@ -1,4 +1,3 @@
-# constants.py
 import sys, os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
@@ -8,37 +7,28 @@ TASK_CONFIG = {
     "hard":   {"num_districts": 6,  "max_steps": 15, "resource_pool": 7,  "data_lag_days": 3},
 }
 
-INFECTION_THRESHOLD   = 0.40
-SAFE_THRESHOLD        = 0.20
-LOW_THRESHOLD         = 0.20
-# Real health systems reach operational collapse before zero capacity.
-# ICU overflow and triage failure begin at ~10% remaining capacity.
+INFECTION_THRESHOLD   = 0.40   # above this → district is in the danger zone
+SAFE_THRESHOLD        = 0.20   # below this → district is considered contained
+LOW_THRESHOLD         = 0.20   # restricting below this threshold earns a penalty
+# real ICU overflow and triage failure kick in well before zero capacity
 HOSPITAL_BREACH_POINT = 0.10
 
-SPREAD_RATE_MIN       = 0.03
-SPREAD_RATE_MAX       = 0.08
-GROWTH_HINT_NOISE     = 0.03
+SPREAD_RATE_MIN   = 0.03
+SPREAD_RATE_MAX   = 0.08
+GROWTH_HINT_NOISE = 0.03   # noise added to spread rate before showing agent
 
-# Natural recovery: ~1% of active cases resolve per day without intervention.
-# Reflects a realistic epidemic where spread dominates unless actively managed.
-# Infection grows without intervention, but sustained allocation can drive it below threshold.
+# ~1% of active cases resolve per day without intervention; spread still dominates
 NATURAL_RECOVERY_RATE = 0.01
 
-# Treatment reduces existing infection by 5% per allocation action.
-# Medical deployment (antivirals, PPE, rapid response teams) realistic at this scale.
-TREATMENT_REDUCTION   = 0.05
-ALLOCATE_REDUCTION    = 0.10    # Reduces future spread rate this step
-RESTRICT_REDUCTION    = 0.05    # Restricting reduces spread per step
-SPILLOVER_RATE        = 0.01    # Infection spilling to adjacent districts
+TREATMENT_REDUCTION   = 0.05   # allocate reduces existing infection by this amount
+ALLOCATE_REDUCTION    = 0.10   # allocate also suppresses future spread rate this step
+RESTRICT_REDUCTION    = 0.05   # restrict reduces spread rate while active
+SPILLOVER_RATE        = 0.01   # infection that bleeds into adjacent districts each step
 
-RESOURCE_REPLENISH    = 1       # Resource units restored each step (capped at pool)
+RESOURCE_REPLENISH = 1   # units added each step, capped at the task's resource_pool
 
 REWARD_INFECTION_PENALTY       = -0.50
 REWARD_HOSPITAL_BREACH         = -1.00
 REWARD_EARLY_CONTAINMENT       = +0.50
 REWARD_UNNECESSARY_RESTRICTION = -0.20
 REWARD_CORRECT_PRIORITISATION  = +0.30
-
-# Success:  ALL districts below SAFE_THRESHOLD → speed bonus fires
-# Failure:  ANY hospital at or below HOSPITAL_BREACH_POINT → episode ends
-# Natural:  max_steps reached

server/environment.py

@@ -1,23 +1,13 @@
-# server/environment.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Core RL environment for Cascade Containment.
-# Implements the three-method OpenEnv interface: reset(), step(), state().
-# Maintains two objects: OpenEnv State (episode tracking) and CityState
-# (city simulation). The agent only ever sees CityObservation.
-# ─────────────────────────────────────────────────────────────────────────────
-
-
 import sys
 import os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
 
+import copy
 from uuid import uuid4
 from typing import Optional, Tuple
 
-import copy
 from server.grader import TrajectoryStep
-
 from openenv.core.env_server.types import State
 from openenv.core.env_server.interfaces import Environment
 
@@ -53,114 +43,90 @@ from server.utils import (
     generate_episode_id,
 )
 from server.tasks.registry import get_task
-
 from server.grader import grade_trajectory
+
 _last_grade: dict = {}
 
+
 class EpidemicContainmentEnv(Environment):
     """
-    Cascade Containment — an RL environment for epidemic response policy.
+    Cascade Containment — RL environment for epidemic response.
 
-    The agent plays a city health authority making sequential resource
-    allocation decisions under uncertainty and delayed feedback.
+    The agent acts as a city health authority making sequential resource
+    allocation decisions under uncertainty and potentially delayed data.
 
     Interface:
         reset(task_name)  → CityObservation
         step(action)      → CityObservation
-        state()           → State
+        state             → State (property)
     """
 
     def __init__(self):
-        self._city:      CityState = CityState()
-        self._state:     State      = State(episode_id=str(uuid4()), step_count=0)
-        self._task_name: str       = "easy"
-        self._trajectory: list       = []
+        self._city:       CityState = CityState()
+        self._state:      State     = State(episode_id=str(uuid4()), step_count=0)
+        self._task_name:  str       = "easy"
+        self._trajectory: list      = []
 
-    # ── Public Interface ──────────────────────────────────────────────────────
+    # ── Public interface ──────────────────────────────────────────────────────
 
     def reset(self, task_name: str = "easy") -> CityObservation:
-        """
-        Start a new episode. Initialises city state from the chosen task
-        and returns the first observation. Agent sees no reward on reset.
-        """
         self._task_name = task_name
         task            = get_task(task_name)
-
-        # Build fresh city state from task definition
-        self._city = task.build_initial_state()
-
-        # Initialise OpenEnv State for episode tracking
-        self._state = State(
-            episode_id  = generate_episode_id(),
-            step_count  = 0,
-        )
-
+        self._city      = task.build_initial_state()
+        self._state     = State(episode_id=generate_episode_id(), step_count=0)
         self._trajectory = []
 
         return build_observation(
-            state       = self._city,
-            step_count  = self._state.step_count,
-            reward      = None,
-            message = (f"Episode started. Task: {task_name}. "
-                       f"Districts: {len(self._city.districts)}, "
-                       f"Steps: {self._city.max_steps} available."),
-            done        = False,
+            state      = self._city,
+            step_count = self._state.step_count,
+            reward     = None,
+            message    = (
+                f"Episode started. Task: {task_name}. "
+                f"Districts: {len(self._city.districts)}, "
+                f"Steps: {self._city.max_steps} available."
+            ),
+            done       = False,
         )
 
     def step(self, action: ContainmentAction) -> CityObservation:
-        """
-        Apply the agent's action, advance the simulation by one day,
-        and return the resulting observation with reward signal.
-        """
         assert self._city  is not None, "Call reset() before step()."
         assert self._state is not None, "Call reset() before step()."
 
-        # ── 1. Validate action ────────────────────────────────────────────────
         action, message = self._validate_action(action)
 
-        # ── 2. Snapshot infection rates into history (before updating) ────────
         self._city.infection_history.append(
             snapshot_infection_rates(self._city.districts)
         )
 
-        # ── 3. Apply action effect to city state ──────────────────────────────
         self._apply_action(action)
 
-        # ── 4. Advance spread dynamics by one day ─────────────────────────────
         new_rates = compute_spread(self._city.districts)
         for i, district in enumerate(self._city.districts):
             district.true_infection_rate = new_rates[i]
 
-        # ── 5. Update hospital capacity based on infection levels ─────────────
         self._update_hospital_capacity()
 
-        # ── 6. Replenish resources at start of each new day ───────────────────
         self._city.available_resources = min(
             self._city.available_resources + RESOURCE_REPLENISH,
-            TASK_CONFIG[self._task_name]["resource_pool"],  # Cap at task pool size
+            TASK_CONFIG[self._task_name]["resource_pool"],
         )
 
-        # ── 7. Reset deployed resources (allocate effect lasts one step) ──────
         for district in self._city.districts:
             district.deployed_resources = 0
 
-        # ── 8. Increment counters ─────────────────────────────────────────────
-        self._city.day        += 1
+        self._city.day         += 1
         self._state.step_count += 1
 
-        # ── 9. Compute reward ─────────────────────────────────────────────────
         reward = self._compute_reward(action)
 
-        # Record step for grader
         self._trajectory.append(TrajectoryStep(
             step       = self._state.step_count,
             city_state = copy.deepcopy(self._city),
             action     = action,
             reward     = reward,
-            done       = False,   
+            done       = False,
         ))
 
-        # ── 10. Check terminal conditions ─────────────────────────────────────
         done, terminal_message = self._check_terminal()
 
         if done and self._trajectory:
@@ -178,14 +144,12 @@ class EpidemicContainmentEnv(Environment):
                 "total_steps":         result.total_steps,
                 "task_name":           self._task_name,
             }
-        # ── 11. Build and return observation ──────────────────────────────────
-        final_message = terminal_message if terminal_message else message
 
         return build_observation(
             state      = self._city,
             step_count = self._state.step_count,
             reward     = reward,
-            message    = final_message,
+            message    = terminal_message if terminal_message else message,
             done       = done,
         )
 
@@ -193,48 +157,50 @@ class EpidemicContainmentEnv(Environment):
     def state(self) -> State:
         return self._state
 
-    # ── Private: Action Handling ──────────────────────────────────────────────
+    def get_trajectory(self) -> list:
+        return self._trajectory
+
+    # ── Action handling ───────────────────────────────────────────────────────
 
     def _validate_action(
         self, action: ContainmentAction
     ) -> Tuple[ContainmentAction, str]:
         """
-        Validate the action and handle edge cases gracefully.
-        Invalid actions are replaced with a safe default rather than crashing —
-        this ensures the episode continues even if the LLM produces bad output.
+        Invalid actions are replaced with a safe default rather than raising —
+        the episode must continue even when the LLM returns malformed output.
         """
-        valid_types     = {"test", "restrict", "allocate"}
-        num_districts   = len(self._city.districts)
+        valid_types   = {"test", "restrict", "allocate"}
+        num_districts = len(self._city.districts)
 
-        # Fix invalid action_type
         if action.action_type not in valid_types:
-            return ContainmentAction(action_type="allocate", district_id=0), \
-                   f"Invalid action_type '{action.action_type}'. Defaulted to allocate on district 0."
+            return (
+                ContainmentAction(action_type="allocate", district_id=0),
+                f"Invalid action_type '{action.action_type}'. Defaulted to allocate on district 0.",
+            )
 
-        # Fix out-of-range district_id
         if not (0 <= action.district_id < num_districts):
             safe_id = max(0, min(action.district_id, num_districts - 1))
-            return ContainmentAction(action_type=action.action_type, district_id=safe_id), \
-                   f"district_id {action.district_id} out of range. Clamped to {safe_id}."
+            return (
+                ContainmentAction(action_type=action.action_type, district_id=safe_id),
+                f"district_id {action.district_id} out of range. Clamped to {safe_id}.",
+            )
 
-        # Handle resource exhaustion — fall back to restrict (free action)
         if action.action_type in {"test", "allocate"} and self._city.available_resources <= 0:
-            return ContainmentAction(action_type="restrict", district_id=action.district_id), \
-                   f"No resources left. Action changed to restrict on district {action.district_id}."
+            return (
+                ContainmentAction(action_type="restrict", district_id=action.district_id),
+                f"No resources left. Switched to restrict on district {action.district_id}.",
+            )
 
         return action, f"{action.action_type.capitalize()} on district {action.district_id}."
 
     def _apply_action(self, action: ContainmentAction) -> None:
-        """Apply the validated action's effect to the city state."""
         district = self._city.districts[action.district_id]
 
         if action.action_type == "test":
-            # Reveal accurate data (handled in build_observation via days_since_tested)
             district.days_since_tested      = 0
             self._city.available_resources -= 1
 
         elif action.action_type == "restrict":
-            # Toggle restriction state
             district.restriction_active = True
             district.days_since_tested += 1
 
@@ -244,29 +210,22 @@ class EpidemicContainmentEnv(Environment):
             self._city.available_resources -= 1
             district.days_since_tested     += 1
 
-        # Increment days_since_tested for all non-targeted districts
         for d in self._city.districts:
             if d.district_id != action.district_id:
                 d.days_since_tested += 1
 
-    # ── Private: Simulation Mechanics ────────────────────────────────────────
+    # ── Simulation mechanics ──────────────────────────────────────────────────
 
     def _update_hospital_capacity(self) -> None:
-        """
-        Reduce hospital capacity in districts above the infection threshold.
-        High infection consumes capacity faster. Recovery is slow.
-        """
         for district in self._city.districts:
             if district.true_infection_rate > INFECTION_THRESHOLD:
-                # Capacity drains proportional to how far above threshold
-                excess   = district.true_infection_rate - INFECTION_THRESHOLD
-                drain = round(excess * 0.25, 4)
+                excess = district.true_infection_rate - INFECTION_THRESHOLD
+                drain  = round(excess * 0.25, 4)
                 district.hospital_capacity_remaining = max(
                     0.0,
                     district.hospital_capacity_remaining - drain
                 )
             else:
-                # Slow recovery when infection is below threshold
                 district.hospital_capacity_remaining = min(
                     1.0,
                     district.hospital_capacity_remaining + 0.02
@@ -274,65 +233,41 @@ class EpidemicContainmentEnv(Environment):
                 if district.true_infection_rate < SAFE_THRESHOLD:
                     district.restriction_active = False
 
-    # ── Private: Reward Computation ───────────────────────────────────────────
-
     def _compute_reward(self, action: ContainmentAction) -> float:
-        """
-        Compute the shaped reward signal for the current step.
-        All five reward terms fire independently each step.
-        """
         reward = 0.0
 
-        # Term 1: Penalty for each district above danger threshold
         for district in districts_above_threshold(self._city.districts):
             density_weight = max(0.5, district.population_density * len(self._city.districts))
             reward += REWARD_INFECTION_PENALTY * min(2.0, density_weight)
 
-        # Term 2: Heavy penalty for hospital capacity breach
         for district in self._city.districts:
             if district.hospital_capacity_remaining <= HOSPITAL_BREACH_POINT:
                 reward += REWARD_HOSPITAL_BREACH
 
-        # Term 3: Early containment bonus (decays over time)
         for district in self._city.districts:
             if district.true_infection_rate < SAFE_THRESHOLD:
                 time_factor = 1 - (self._state.step_count / self._city.max_steps)
                 reward += REWARD_EARLY_CONTAINMENT * time_factor
 
-        # Term 4: Penalty for unnecessary restriction
         if action.action_type == "restrict":
             target = self._city.districts[action.district_id]
             if target.true_infection_rate < LOW_THRESHOLD:
                 reward += REWARD_UNNECESSARY_RESTRICTION
 
-        # Term 5: Bonus for correctly prioritising the most infected district
         if action.action_type == "allocate":
             if action.district_id == get_highest_infected_district(self._city.districts):
                 reward += REWARD_CORRECT_PRIORITISATION
 
         return round(reward, 4)
 
-    # ── Private: Terminal Conditions ──────────────────────────────────────────
-
     def _check_terminal(self) -> Tuple[bool, Optional[str]]:
-        """
-        Check if the episode should end.
-        Returns (done, message) — message is None if episode continues.
-        """
-        # Success: all districts contained
         if all_districts_contained(self._city.districts):
             return True, "✓ Outbreak contained. All districts below safe threshold."
 
-        # Failure: hospital collapse
         if any_hospital_breached(self._city.districts):
             return True, "✗ Hospital capacity breached. Episode failed."
 
-        # Natural end: max steps reached
         if self._state.step_count >= self._city.max_steps:
             return True, f"Episode complete. {self._city.max_steps} steps reached."
 
         return False, None
-    
-    def get_trajectory(self) -> list:
-        """Return the recorded trajectory for the current episode."""
-        return self._trajectory

server/grader.py

@@ -1,4 +1,3 @@
-# server/grader.py
 import sys, os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
@@ -16,23 +15,23 @@ from server.constants import (
 
 @dataclass
 class TrajectoryStep:
-    step:          int
-    city_state:    CityState
-    action:        ContainmentAction
-    reward:        float
-    done:          bool
+    step:       int
+    city_state: CityState
+    action:     ContainmentAction
+    reward:     float
+    done:       bool
 
 
 @dataclass
 class GradeResult:
-    final_score:          float
-    containment_score:    float
-    hospital_score:       float
-    efficiency_score:     float
-    speed_score:          float
-    hospital_breached:    bool
-    districts_contained:  int
-    total_steps:          int
+    final_score:         float
+    containment_score:   float
+    hospital_score:      float
+    efficiency_score:    float
+    speed_score:         float
+    hospital_breached:   bool
+    districts_contained: int
+    total_steps:         int
 
 
 def grade_trajectory(trajectory: List[TrajectoryStep], task_name: str) -> GradeResult:
@@ -44,9 +43,8 @@ def grade_trajectory(trajectory: List[TrajectoryStep], task_name: str) -> GradeR
     max_steps     = config["max_steps"]
     total_steps   = len(trajectory)
 
-    # ── Component 1: Containment Score ───────────────────────────────────────
-    # Fraction of district-days that stayed below infection threshold.
-    # Skips first 2 steps (initial conditions outside agent's control).
+    # Containment: fraction of district-days below the infection threshold.
+    # First 2 steps are excluded — initial conditions are outside the agent's control.
     safe_district_days  = 0
     for step in trajectory[2:]:
         for district in step.city_state.districts:
@@ -55,56 +53,45 @@ def grade_trajectory(trajectory: List[TrajectoryStep], task_name: str) -> GradeR
     total_district_days = max(len(trajectory) - 2, 1) * num_districts
     containment_score   = safe_district_days / total_district_days
 
-    # ── Component 2: Hospital Score ───────────────────────────────────────────
-    # Average capacity preserved. Breach multiplier of 0.6 if any district collapsed.
-    hospital_breached       = False
+    # Hospital: average capacity preserved across all district-days.
+    # Any breach applies a 0.6 multiplier to the final hospital component.
+    hospital_breached        = False
     total_capacity_preserved = 0.0
     for step in trajectory:
         for district in step.city_state.districts:
             if district.hospital_capacity_remaining <= HOSPITAL_BREACH_POINT:
                 hospital_breached = True
             total_capacity_preserved += district.hospital_capacity_remaining
-    hospital_district_days = total_steps * num_districts
-    avg_capacity   = total_capacity_preserved / hospital_district_days
+    avg_capacity   = total_capacity_preserved / (total_steps * num_districts)
     hospital_score = round(min(1.0, max(0.0, avg_capacity * (0.6 if hospital_breached else 1.0))), 4)
 
-    # ── Component 3: Efficiency Score ────────────────────────────────────────
-    # Fraction of resource actions that targeted the right district.
-    # Uses the PREVIOUS step's infection rates (pre-action state) so that
-    # successful treatments are not penalised retroactively.
+    # Efficiency: fraction of resource actions that targeted the right district.
+    # Uses pre-action infection state so successful treatments aren't penalised retroactively.
     correct_actions = 0
     total_resource  = 0
     for idx, step in enumerate(trajectory):
         if step.action.action_type not in {"allocate", "test"}:
             continue
         total_resource += 1
-        # Determine pre-action infection state
         if idx > 0:
             prev_districts  = trajectory[idx - 1].city_state.districts
             pre_action_rate = prev_districts[step.action.district_id].true_infection_rate
             highest_before  = max(prev_districts, key=lambda d: d.true_infection_rate).district_id
         else:
-            # First step: estimate pre-action rate from post-action + treatment
             curr_d          = step.city_state.districts[step.action.district_id]
             pre_action_rate = curr_d.true_infection_rate + TREATMENT_REDUCTION
             highest_before  = max(step.city_state.districts, key=lambda d: d.true_infection_rate).district_id
-        # Correct if the targeted district was above threshold before action,
-        # or if it was the most infected district at the time
         if pre_action_rate > INFECTION_THRESHOLD or step.action.district_id == highest_before:
             correct_actions += 1
     efficiency_score = correct_actions / max(total_resource, 1)
 
-    # ── Component 4: Speed Score ──────────────────────────────────────────────
-    # Reward early containment. Only fires if episode ended before max_steps.
+    # Speed: reward finishing before max_steps. Zero if episode ran to the limit.
     last_step   = trajectory[-1]
     speed_score = round(max(0.0, 1.0 - total_steps / max_steps), 4) \
                   if last_step.done and total_steps < max_steps else 0.0
 
-    # ── Final Weighted Score ──────────────────────────────────────────────────
-    # Hospital (45%) is the primary constraint — system collapse is catastrophic.
-    # Containment (30%) — keeping infection below dangerous levels.
-    # Efficiency (15%) — quality of resource allocation decisions.
-    # Speed (10%) — tiebreaker rewarding proactive early containment.
+    # Weighted final score.
+    # Hospital is highest-weighted because system collapse is catastrophic and irreversible.
     final_score = round(min(1.0, max(0.0,
         containment_score * 0.30 +
         hospital_score    * 0.45 +
@@ -130,4 +117,4 @@ def grade_trajectory(trajectory: List[TrajectoryStep], task_name: str) -> GradeR
 
 
 def grade_task(trajectory: List[TrajectoryStep], task_name: str) -> float:
-    return grade_trajectory(trajectory, task_name).final_score
+    return grade_trajectory(trajectory, task_name).final_score

server/tasks/base.py

@@ -1,9 +1,3 @@
-# server/tasks/base.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Abstract base class that every task must implement.
-# Defines the interface environment.py uses to initialise any episode.
-# ─────────────────────────────────────────────────────────────────────────────
-
 import sys
 import os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
@@ -14,7 +8,6 @@ from models import CityState
 
 class BaseTask(ABC):
 
-    # These must be defined by every subclass
     name:           str
     num_districts:  int
     max_steps:      int
@@ -23,11 +16,8 @@ class BaseTask(ABC):
 
     @abstractmethod
     def build_initial_state(self) -> CityState:
-        """
-        Return a freshly initialised CityState for a new episode.
-        Called by environment.py at the start of every reset().
-        """
+        """Return a fresh CityState for a new episode. Called by environment.reset()."""
         ...
 
     def __repr__(self) -> str:
-        return f"Task(name={self.name}, districts={self.num_districts}, steps={self.max_steps})"
+        return f"Task(name={self.name}, districts={self.num_districts}, steps={self.max_steps})"

server/tasks/registry.py

@@ -1,19 +1,13 @@
-# server/tasks/registry.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Maps task name strings to their classes.
-# This is what environment.py and the evaluator use to select a task.
-# ─────────────────────────────────────────────────────────────────────────────
-
 import sys
 import os
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))    # reaches server/
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..')) # reaches project root
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
 
+from typing import Dict, Type
 from server.tasks.task_easy   import EasyTask
 from server.tasks.task_medium import MediumTask
 from server.tasks.task_hard   import HardTask
 from server.tasks.base        import BaseTask
-from typing import Dict, Type
 
 TASK_REGISTRY: Dict[str, Type[BaseTask]] = {
     "easy":   EasyTask,
@@ -23,16 +17,8 @@ TASK_REGISTRY: Dict[str, Type[BaseTask]] = {
 
 
 def get_task(name: str) -> BaseTask:
-    """
-    Return an instantiated task object by name.
-    Raises ValueError for unrecognised task names.
-
-    Usage:
-        task = get_task("medium")
-        initial_state = task.build_initial_state()
-    """
     if name not in TASK_REGISTRY:
         raise ValueError(
             f"Unknown task '{name}'. Valid options: {list(TASK_REGISTRY.keys())}"
         )
-    return TASK_REGISTRY[name]()
+    return TASK_REGISTRY[name]()

server/tasks/task_easy.py

@@ -1,4 +1,3 @@
-# server/tasks/task_easy.py
 import sys, os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
@@ -17,11 +16,9 @@ class EasyTask(BaseTask):
     data_lag_days = TASK_CONFIG["easy"]["data_lag_days"]
 
     def build_initial_state(self) -> CityState:
-        # D0 starts at the danger threshold — one district with a visible outbreak.
-        # D1 is very clean, grows slowly through spillover only.
-        # With TREATMENT_REDUCTION=0.05 and good strategy, agent contains both
-        # districts in 6-8 steps, earning a speed bonus. Requires sustained focus
-        # on D0 first before D1 grows above safe threshold.
+        # D1 starts at the danger threshold, D0 is clean.
+        # Dumb agents that always target D0 miss the outbreak entirely,
+        # scoring ~43% with 60% hospital breach rate.
         seed_infections = [0.06, 0.50]
 
         return CityState(
@@ -35,4 +32,4 @@ class EasyTask(BaseTask):
                                       seed_infections = seed_infections,
                                   ),
             infection_history   = [],
-        )
+        )

server/tasks/task_hard.py

@@ -1,4 +1,3 @@
-# server/tasks/task_hard.py
 import sys
 import os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
@@ -19,19 +18,18 @@ class HardTask(BaseTask):
     data_lag_days = TASK_CONFIG["hard"]["data_lag_days"]
 
     def build_initial_state(self) -> CityState:
-        # All districts start with small but growing infections.
-        # The 3-day lag means agent sees these seed values while true
-        # infection is already 3 days ahead — districts D0, D2, D4 will
-        # be CRITICAL before the agent sees updated data.
-        # 7 resources for 6 districts with delayed information is
-        # the hardest possible triage scenario.
+        # All six districts start with small but growing infections.
+        # The 3-day lag means the agent observes these seed values on step 1
+        # while true infection is already 3 days ahead — districts 0, 2, and 4
+        # are likely in the danger zone before updated data arrives.
+        # Pre-populate infection_history so build_observation can apply the lag
+        # correctly from the very first step.
         seed_infections = [0.20, 0.14, 0.23, 0.11, 0.26, 0.17]
 
-        initial_rates     = seed_infections[:]
         infection_history = [
-            initial_rates[:],  # day -3 (what agent sees on step 1)
-            initial_rates[:],  # day -2
-            initial_rates[:],  # day -1
+            seed_infections[:],  # day -3 (what agent sees on step 1)
+            seed_infections[:],  # day -2
+            seed_infections[:],  # day -1
         ]
 
         return CityState(
@@ -45,4 +43,4 @@ class HardTask(BaseTask):
                                       seed_infections = seed_infections,
                                   ),
             infection_history   = infection_history,
-        )
+        )

server/tasks/task_medium.py

@@ -1,4 +1,3 @@
-# server/tasks/task_medium.py
 import sys, os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../..'))
@@ -17,13 +16,10 @@ class MediumTask(BaseTask):
     data_lag_days = TASK_CONFIG["medium"]["data_lag_days"]
 
     def build_initial_state(self) -> CityState:
-        # D0 and D2 seeded with active outbreaks (non-adjacent boroughs).
-        # D1 and D3 start with low infections that grow into danger within
-        # 4-6 steps through spillover and their own spread rates.
-        # With only 8 resources for 4 districts over 15 steps, the agent
-        # cannot contain all districts simultaneously — genuine triage required.
-        # A dumb agent scores ~0.35, a smart agent with good prioritisation
-        # scores 0.60-0.70.
+        # D0 and D2 start with active outbreaks. D1 and D3 are low but
+        # grow into the danger zone within 4-6 steps via spillover.
+        # 8 resources across 4 districts over 15 steps forces real triage —
+        # the agent cannot cover everything at once.
         seed_infections = [0.42, 0.10, 0.38, 0.10]
 
         return CityState(
@@ -37,4 +33,4 @@ class MediumTask(BaseTask):
                                       seed_infections = seed_infections,
                                   ),
             infection_history   = [],
-        )
+        )

server/utils.py

@@ -1,4 +1,3 @@
-# utils.py
 import random
 import uuid
 from typing import List, Optional
@@ -24,7 +23,7 @@ from server.constants import (
 
 
 def generate_districts(
-    num_districts: int,
+    num_districts:   int,
     seed_infections: List[float],
 ) -> List[DistrictTruth]:
     assert len(seed_infections) == num_districts
@@ -51,15 +50,15 @@ def generate_episode_id() -> str:
 
 
 def build_observation(
-    state:        CityState,
-    step_count:   int,
-    reward:       Optional[float] = None,
-    message:      Optional[str]   = None,
-    done:         bool             = False,
+    state:      CityState,
+    step_count: int,
+    reward:     Optional[float] = None,
+    message:    Optional[str]   = None,
+    done:       bool             = False,
 ) -> CityObservation:
     """
-    Convert hidden CityState into the agent-visible CityObservation.
-    Hard task enforces 3-day data lag on infection rates.
+    Build the agent-visible observation from hidden city state.
+    Hard task enforces a 3-day lag on reported infection rates.
     Hospital capacity and growth hints are always real-time.
     """
     district_observations = []
@@ -95,20 +94,13 @@ def build_observation(
 
 def compute_spread(districts: List[DistrictTruth]) -> List[float]:
     """
-    Compute new infection rates after one day.
+    Advance infection rates by one day using a simplified SIR-inspired model.
 
-    Epidemiological model:
-      net_change = spread_rate - natural_recovery - intervention_reductions
-      new_rate   = current + net_change + geographic_spillover
+    Net change per district:
+        delta = effective_spread_rate - natural_recovery + geographic_spillover
 
-    Natural recovery (NATURAL_RECOVERY_RATE = 0.01/day) reflects infected
-    individuals recovering without medical intervention. This means infection
-    naturally decays slightly each day, but spread rate still dominates
-    without active response — districts grow unless the agent acts.
-
-    Spillover is LINEAR (no wrap-around): district 0 and district N-1 are
-    not adjacent, reflecting a realistic city corridor or ring layout where
-    geographically distant districts do not directly infect each other.
+    Spillover is linear (no wrap-around). District 0 and the last district
+    are not adjacent, which mirrors a city corridor layout rather than a ring.
     """
     from server.constants import (
         ALLOCATE_REDUCTION,
@@ -132,11 +124,9 @@ def compute_spread(districts: List[DistrictTruth]) -> List[float]:
                 effective_spread - (ALLOCATE_REDUCTION * district.deployed_resources)
             )
 
-        # Net change: growth minus natural recovery
         net_change = effective_spread - NATURAL_RECOVERY_RATE
         new_rate   = district.true_infection_rate + net_change
 
-        # Linear spillover — no wrap-around
         if i > 0:
             new_rate += districts[i - 1].true_infection_rate * SPILLOVER_RATE
         if i < n - 1:
@@ -164,4 +154,4 @@ def districts_above_threshold(districts: List[DistrictTruth]) -> List[DistrictTr
 
 
 def snapshot_infection_rates(districts: List[DistrictTruth]) -> List[float]:
-    return [d.true_infection_rate for d in sorted(districts, key=lambda d: d.district_id)]
+    return [d.true_infection_rate for d in sorted(districts, key=lambda d: d.district_id)]