Improve agent: fix strategy rule, chain-of-thought, 8 rollouts, better memory

baseline/evaluator.py

@@ -1,17 +1,8 @@
 # baseline/evaluator.py
-# ─────────────────────────────────────────────────────────────────────────────
-# GRPO-style evaluation loop for Cascade Containment.
-# Imports core components — stays focused on orchestration only.
-# ─────────────────────────────────────────────────────────────────────────────
-
-import os
-import sys
+import os, sys, time
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
-
-import time
-from typing import List, Tuple
+from typing import List, Tuple, Any
 from openai import OpenAI
-from typing import Any
 
 from client import CascadeContainmentEnv
 from models import ContainmentAction, CityObservation
@@ -19,16 +10,12 @@ from baseline.policy import get_client, build_prompt, call_llm, parse_action
 from core.trajectory import EpisodicMemory
 from core.reward import normalise_score
 from core.policy_update import compute_advantage, update_memory
-
 import requests as http_requests
 
-N_ROLLOUTS = 5
+N_ROLLOUTS = 8   # increased from 5 — GRPO needs more rollouts for stable learning signal
 
 
-# ── Prompt Builder With Memory ────────────────────────────────────────────────
-
 def build_prompt_with_memory(obs: CityObservation, memory: EpisodicMemory) -> str:
-    """Extend base prompt with retrieved memories from similar past situations."""
     from baseline.policy import build_prompt
     base         = build_prompt(obs)
     memory_block = memory.retrieve(obs)
@@ -36,25 +23,15 @@ def build_prompt_with_memory(obs: CityObservation, memory: EpisodicMemory) -> st
     if not memory_block:
         return base
 
-    injection = (
-        "\n"
-        + memory_block
-        + "\nUse these past experiences to make a better decision.\n"
-    )
-    return base.replace("Your decision:", injection + "Your decision:")
-
+    injection = "\n" + memory_block + "\nApply these lessons to your current decision.\n"
+    return base.replace("Your response:", injection + "Your response:")
 
-# ── Single Rollout ────────────────────────────────────────────────────────────
 
 def run_rollout(
-    env:       Any,
-    task_name: str,
-    client:    OpenAI,
-    memory:    EpisodicMemory,
-    verbose:   bool = True,
+    env: Any, task_name: str, client: OpenAI,
+    memory: EpisodicMemory, verbose: bool = True,
 ) -> Tuple[float, int, List[dict]]:
-    """Run one complete episode using memory-augmented prompts."""
-    result = env.reset(task_name=task_name)
+    result       = env.reset(task_name=task_name)
     obs          = result.observation
     done         = result.done
     total_reward = 0.0
@@ -67,11 +44,11 @@ def run_rollout(
         action   = parse_action(response, len(obs.districts))
 
         try:
-            result   = env.step(action)
+            result = env.step(action)
         except Exception as e:
             if "close frame" in str(e).lower() or "websocket" in str(e).lower():
                 if verbose:
-                    print(f"      ⚠ WebSocket dropped at step {step+1}, ending rollout early")
+                    print(f"      ⚠ WebSocket dropped at step {step+1}, ending early")
                 break
             raise
 
@@ -81,17 +58,12 @@ def run_rollout(
         total_reward += reward
         step         += 1
 
-        trajectory.append({
-            "obs":    obs,
-            "action": action,
-            "reward": reward,
-        })
+        trajectory.append({"obs": obs, "action": action, "reward": reward})
 
         if verbose:
             print(
                 f"      step {step:2d}: {action.action_type:8} "
-                f"→ district {action.district_id} "
-                f"| reward: {reward:+.4f}"
+                f"→ district {action.district_id} | reward: {reward:+.4f}"
             )
 
         obs = next_obs
@@ -101,16 +73,10 @@ def run_rollout(
     return total_reward, step, trajectory
 
 
-# ── GRPO Task Runner ──────────────────────────────────────────────────────────
-
 def run_task_grpo(
-    env:       Any,
-    task_name: str,
-    client:    OpenAI,
-    base_url:  str,
-    verbose:   bool = True,
+    env: Any, task_name: str, client: OpenAI,
+    base_url: str, verbose: bool = True,
 ) -> float:
-    """GRPO-style simulated learning loop for one task."""
     if verbose:
         print(f"\n  Task: {task_name.upper()} | {N_ROLLOUTS} rollouts")
         print(f"  {'─'*44}")
@@ -123,14 +89,13 @@ def run_task_grpo(
             label = "base prompt" if len(memory) == 0 else f"memory: {len(memory)} entries"
             print(f"\n    Rollout {i+1}/{N_ROLLOUTS} [{label}]")
 
-        total_reward, steps, trajectory = run_rollout(env, task_name, client, memory, verbose)
+        total_reward, steps, trajectory = run_rollout(
+            env, task_name, client, memory, verbose
+        )
 
-        # Get proper grader score from server
         num_districts = {"easy": 2, "medium": 4, "hard": 6}.get(task_name, 2)
         try:
-            grade_resp = http_requests.get(
-                base_url.rstrip('/') + '/grade', timeout=10
-            )
+            grade_resp = http_requests.get(base_url.rstrip('/') + '/grade', timeout=10)
             if grade_resp.status_code == 200:
                 data  = grade_resp.json()
                 score = data["final_score"]
@@ -146,26 +111,24 @@ def run_task_grpo(
         except Exception:
             score = normalise_score(total_reward, steps, num_districts)
 
-        # Append BEFORE advantage computation
         rollouts.append((total_reward, steps, score))
 
         if verbose:
             print(f"    → Reward: {total_reward:+.4f} | Score: {score:.4f}")
 
-        # ── GRPO advantage computation and memory update ──────────────────────
         completed_rewards = [r[0] for r in rollouts]
         advantage         = compute_advantage(total_reward, completed_rewards[:-1])
         stored            = update_memory(memory, trajectory, advantage)
 
         if verbose:
             mean = sum(completed_rewards[:-1]) / max(len(completed_rewards) - 1, 1) \
-                if len(completed_rewards) > 1 else total_reward
+                   if len(completed_rewards) > 1 else total_reward
             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)
-    best_score = max(rollouts, key=lambda x: x[2])[2]
+    best_score  = max(rollouts, key=lambda x: x[2])[2]
 
     if verbose:
         print(f"\n  Rewards:    {[round(r, 4) for r in all_rewards]}")
@@ -176,13 +139,7 @@ def run_task_grpo(
     return best_score
 
 
-# ── Full Evaluator ────────────────────────────────────────────────────────────
-
-def run_evaluation(
-    base_url: str  = "http://localhost:7860",
-    verbose:  bool = True,
-) -> dict:
-    """Run all three tasks with GRPO episodic memory learning."""
+def run_evaluation(base_url: str = "http://localhost:7860", verbose: bool = True) -> dict:
     if verbose:
         print("\n" + "="*52)
         print("  CASCADE CONTAINMENT — GRPO EVALUATION")
@@ -209,10 +166,8 @@ def run_evaluation(
                     traceback.print_exc()
 
     scores["average"] = round(
-        sum(v for k, v in scores.items() if k != "average") / 3,
-        4
+        sum(v for k, v in scores.items() if k != "average") / 3, 4
     )
-
     elapsed = round(time.time() - start, 1)
 
     if verbose:

baseline/policy.py

@@ -1,10 +1,6 @@
 # baseline/policy.py
-import os
-import sys
+import os, sys, json, re
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
-
-import json
-import re
 from openai import OpenAI
 from models import CityObservation, ContainmentAction
 
@@ -27,21 +23,24 @@ def build_prompt(obs: CityObservation) -> str:
     )
 
     lines = [
-        "You are an epidemic response coordinator.",
-        "Your goal: reduce infection in all districts and prevent hospital collapse.",
+        "You are an epidemic response coordinator making life-or-death resource decisions.",
+        "Your goal: keep infection BELOW 0.40 in all districts and prevent hospital collapse.",
         "",
         f"Step {obs.current_step}/{obs.max_steps} | Resources remaining: {obs.available_resources}",
-        "ℹ️  Resources replenish by 1 each step but cannot exceed your starting pool.",
-        "   Spend wisely — you can never have more resources than the starting amount.",
+        "ℹ️  Resources replenish by 1 each step, capped at starting pool. Spend carefully.",
         "",
-        "Districts (sorted by infection rate, highest first):",
+        "Districts (sorted by CURRENT infection, highest threat first):",
     ]
 
     for d in sorted_districts:
         if d.reported_infection_rate > 0.4:
             status = "🔴 CRITICAL"
         elif d.reported_infection_rate > 0.2:
-            status = "🟡 WARNING"
+            # Add escalation warning based on growth hint
+            if d.growth_rate_hint > 0.06:
+                status = "🟡 WARNING→CRITICAL SOON"
+            else:
+                status = "🟡 WARNING"
         else:
             status = "🟢 SAFE"
 
@@ -52,7 +51,7 @@ def build_prompt(obs: CityObservation) -> str:
         else:
             hosp_status = "hospital OK"
 
-        lag_note = " [DATA IS 3 DAYS OLD]" if has_data_lag else ""
+        lag_note = " [3-DAY OLD DATA]" if has_data_lag else ""
         lines.append(
             f"  D{d.district_id}: {status} infection={d.reported_infection_rate:.2f}{lag_note} "
             f"growth={d.growth_rate_hint:.2f} {hosp_status}({d.hospital_capacity_remaining:.2f})"
@@ -63,45 +62,56 @@ def build_prompt(obs: CityObservation) -> str:
     if not has_data_lag:
         lines += [
             "HOW ACTIONS WORK:",
-            "  - 'allocate': costs 1 resource. REDUCES existing infection AND slows spread.",
-            "    Infections naturally recover slightly each day, but spread dominates",
-            "    without intervention. Sustained allocation is needed to drive below safe.",
-            "  - 'restrict': FREE. Slows future spread. Does NOT reduce existing infection.",
-            "    Use when resources are low or for WARNING districts.",
-            "  - 'test': costs 1 resource. Reveals accurate data. (Not needed — data is real-time.)",
+            "  - 'allocate': costs 1 resource. REDUCES existing infection by 5% AND slows spread.",
+            "    Infections naturally recover 1%/day but spread (3-8%/day) dominates without action.",
+            "    You need SUSTAINED allocation (multiple steps) to drive a district below safe level.",
+            "  - 'restrict': FREE. Slows future spread only. Does NOT reduce existing infection.",
+            "    Use only when you have no resources OR for districts already below 0.20.",
+            "  - 'test': wastes 1 resource. Data is already real-time. NEVER use this.",
+            "",
+            "DECISION RULES — follow this priority order every step:",
+            "1. HOSPITAL EMERGENCY: If ANY hospital < 0.30 capacity → allocate on that district NOW.",
+            "2. TRIAGE: Look at ALL districts. Find the one with the HIGHEST infection rate right now.",
+            "   That is your target this step. Do not stick to the same district if another is worse.",
+            "3. CRITICAL DISTRICT (above 0.40 and have resources): allocate on the highest.",
+            "4. WARNING DISTRICT (0.20-0.40) with growth > 0.06 AND resources available:",
+            "   allocate on it NOW to prevent it from becoming CRITICAL next step.",
+            "5. If resources = 0: restrict on the highest infected district.",
+            "6. NEVER restrict a district below 0.20 — you will be penalised.",
+            "7. NEVER use 'test' — it wastes a resource you cannot afford.",
             "",
-            "STRATEGY — follow in order:",
-            "1. If ANY hospital is below 0.3 capacity: 'allocate' on that district IMMEDIATELY.",
-            "2. Find the district with HIGHEST infection rate.",
-            "3. If it is above 0.2 and you have resources: 'allocate' on it.",
-            "4. Keep allocating to the SAME district next step too.",
-            "   Only switch when that district drops below 0.2 (safe).",
-            "5. If resources = 0: 'restrict' on the highest infected district.",
-            "6. NEVER use 'test' — data is real-time and accurate.",
-            "7. NEVER restrict a district below 0.2 — you will be penalised.",
+            "KEY INSIGHT: Infection spreads 3-8% per day. A WARNING district at 0.38 with",
+            "growth=0.07 will be CRITICAL next step. Act before it escalates, not after.",
         ]
     else:
         lines += [
             "HOW ACTIONS WORK:",
             "  - 'allocate': costs 1 resource. Reduces infection AND slows spread.",
-            "    Data is 3 days old — act on growth_hint to anticipate true state.",
             "  - 'restrict': FREE. Slows future spread only.",
-            "  - 'test': costs 1 resource. (NOT recommended — lag is unavoidable.)",
+            "  - Data is 3 DAYS OLD — you cannot see current true infection rates.",
+            "    Use growth_hint to estimate which districts are getting worse fastest.",
             "",
-            "STRATEGY — follow in order:",
-            "1. If ANY hospital is below 0.3 capacity: 'allocate' on that district IMMEDIATELY.",
-            "2. If resources > 0: 'allocate' on the district with HIGHEST growth_hint.",
-            "   High growth_hint means true infection is accelerating — act early.",
-            "3. If resources = 0: 'restrict' on the district with highest growth_hint.",
-            "4. NEVER use 'test' — spending resources on information wastes your budget.",
+            "DECISION RULES for delayed-information scenario:",
+            "1. HOSPITAL EMERGENCY: If ANY hospital < 0.30 → allocate on that district NOW.",
+            "   Hospital capacity IS real-time even when infection data is lagged.",
+            "2. TRIAGE under uncertainty: Combine lagged infection + growth_hint to estimate severity.",
+            "   A district with infection=0.20 (3 days ago) and growth=0.08 is NOW likely at ~0.44.",
+            "   Formula: estimated_current = reported_infection + 3 × growth_hint",
+            "3. Allocate on the district with HIGHEST estimated current infection.",
+            "4. If resources = 0: restrict on the district with highest growth_hint.",
+            "5. NEVER use 'test' — the 3-day lag is structural, testing does not help.",
         ]
 
     lines += [
         "",
-        "Respond with ONLY valid JSON. No explanation. Example:",
+        "Think briefly (1 sentence): Which district is most dangerous RIGHT NOW and why?",
+        "Then give your JSON decision.",
+        "",
+        "Example response:",
+        'District 0 is critical at 0.65 and growing fastest.',
         '{"action_type": "allocate", "district_id": 0}',
         "",
-        "Your decision:",
+        "Your response:",
     ]
 
     return "\n".join(lines)
@@ -113,14 +123,15 @@ def call_llm(prompt: str, client: OpenAI) -> str:
         messages = [
             {
                 "role":    "system",
-                "content": "You are an epidemic response AI. Always respond with valid JSON only. No explanation, no markdown."
+                "content": (
+                    "You are an epidemic response AI. "
+                    "First write one sentence of reasoning, then a JSON action on the next line. "
+                    "JSON must be valid and contain action_type and district_id."
+                )
             },
-            {
-                "role":    "user",
-                "content": prompt
-            }
+            {"role": "user", "content": prompt}
         ],
-        max_tokens  = 50,
+        max_tokens  = 120,   # increased to allow brief reasoning + JSON
         temperature = 0.1,
     )
     return (response.choices[0].message.content or "").strip()

core/trajectory.py

@@ -1,14 +1,6 @@
 # core/trajectory.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Episodic memory for GRPO-style simulated learning.
-# Stores high-reward (observation, action, reward) tuples from past rollouts.
-# Retrieved at each step to provide instance-level guidance to the policy.
-# ─────────────────────────────────────────────────────────────────────────────
-
-import os
-import sys
+import os, sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
-
 from typing import List
 from models import ContainmentAction, CityObservation
 
@@ -16,11 +8,10 @@ from models import ContainmentAction, CityObservation
 class EpisodicMemory:
     """
     Stores high-reward steps from past rollouts.
-    Retrieved by similarity to current observation to guide next rollout.
-
-    This is the contextual bandit component — the agent gets specific examples:
-    "last time infection was [0.45, 0.12] with 7 resources,
-     allocating to district 0 earned +0.3"
+    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.
     """
 
     def __init__(self, max_size: int = 20):
@@ -28,23 +19,31 @@ class EpisodicMemory:
         self.max_size             = max_size
 
     def store(self, obs: CityObservation, action: ContainmentAction, reward: float):
-        """Store a step only if it earned positive reward."""
-        if reward < -0.3:
+        """Store a step only if it earned meaningful positive reward."""
+        if reward < 0.0:   # stricter threshold — only store clearly positive steps
             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"
+
         self.memories.append({
             "infection_profile": [round(d.reported_infection_rate, 2) for d in obs.districts],
             "resources":         obs.available_resources,
+            "phase":             phase,
             "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),
         })
 
         # 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 = 3) -> str:
+    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.
@@ -54,27 +53,30 @@ class EpisodicMemory:
             return ""
 
         current = [round(d.reported_infection_rate, 2) for d in obs.districts]
+        phase   = "early" if obs.current_step <= obs.max_steps // 3 else \
+                  "mid"   if obs.current_step <= 2 * obs.max_steps // 3 else "late"
 
-        def l1_distance(memory: dict) -> float:
+        def score(memory: dict) -> float:
             profile = memory["infection_profile"]
             if len(profile) != len(current):
                 return float("inf")
-            return sum(abs(a - b) for a, b in zip(profile, current))
+            l1 = sum(abs(a - b) for a, b in zip(profile, current))
+            # Slight preference for matching episode phase
+            phase_bonus = 0.0 if memory.get("phase") == phase else 0.1
+            return l1 + phase_bonus
 
-        ranked = sorted(self.memories, key=l1_distance)
+        ranked = sorted(self.memories, key=score)
         top    = ranked[:top_k]
 
-        lines  = ["Relevant past decisions (from successful rollouts):"]
+        lines = ["Past decisions that earned positive reward (use as guidance):"]
         for m in top:
             lines.append(
-                f"  - Profile {m['infection_profile']} | resources={m['resources']}: "
-                f"'{m['action_type']}' on district {m['district_id']} "
-                f"→ reward {m['reward']:+.4f}"
+                f"  Phase={m.get('phase','?')} Profile={m['infection_profile']} resources={m['resources']}: "
+                f"'{m['action_type']}' D{m['district_id']} → reward {m['reward']:+.2f}"
             )
         return "\n".join(lines)
 
     def clear(self):
-        """Clear memory between tasks — memories are task-specific."""
         self.memories = []
 
     def __len__(self) -> int: