Fix resource replenishment, grader efficiency, seed calibration, prompt strategy

baseline/policy.py

@@ -1,10 +1,4 @@
 # baseline/policy.py
-# ─────────────────────────────────────────────────────────────────────────────
-# LLM-based policy for Cascade Containment.
-# Reads CityObservation, calls LLM via OpenAI client, returns ContainmentAction.
-# Uses environment variables for API configuration as required by hackathon rules.
-# ─────────────────────────────────────────────────────────────────────────────
-
 import os
 import sys
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
@@ -15,64 +9,91 @@ from openai import OpenAI
 from models import CityObservation, ContainmentAction
 
 
-# ── Client Setup ──────────────────────────────────────────────────────────────
-
 def get_client() -> OpenAI:
-    """
-    Initialise OpenAI client from environment variables.
-    Required by hackathon rules — never hardcode API keys.
-    """
     return OpenAI(
-        api_key  = os.environ.get("HF_TOKEN", ""),
+        api_key  = os.environ.get("API_KEY") or os.environ.get("HF_TOKEN", ""),
         base_url = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1"),
     )
 
 
-# ── Prompt Builder ────────────────────────────────────────────────────────────
-
 def build_prompt(obs: CityObservation) -> str:
-    """
-    Convert a CityObservation into a clear, structured prompt.
-    The prompt gives the LLM everything it needs to make an informed decision.
-    """
+    num_districts = len(obs.districts)
+    has_data_lag  = num_districts == 6  # only hard task has lagged data
+
+    sorted_districts = sorted(
+        obs.districts,
+        key=lambda d: d.reported_infection_rate,
+        reverse=True
+    )
+
     lines = [
-        "You are a public health authority managing an epidemic outbreak.",
-        "Your goal is to contain infection across all districts before hospitals collapse.",
+        "You are an epidemic response coordinator.",
+        "Your goal: reduce infection in all districts and prevent hospital collapse.",
         "",
-        f"Current situation (Step {obs.current_step}/{obs.max_steps}):",
-        f"Available resources: {obs.available_resources}",
+        f"Step {obs.current_step}/{obs.max_steps} | Resources remaining: {obs.available_resources}",
+        "⚠️  Resources do NOT replenish. Every resource spent is gone permanently.",
         "",
-        "District status:",
+        "Districts (sorted by infection rate, highest first):",
     ]
 
-    for d in obs.districts:
-        status = "DANGER" if d.reported_infection_rate > 0.4 else \
-                 "WARNING" if d.reported_infection_rate > 0.2 else "SAFE"
+    for d in sorted_districts:
+        if d.reported_infection_rate > 0.4:
+            status = "🔴 CRITICAL"
+        elif d.reported_infection_rate > 0.2:
+            status = "🟡 WARNING"
+        else:
+            status = "🟢 SAFE"
+
+        if d.hospital_capacity_remaining < 0.3:
+            hosp_status = "⚠️ HOSPITAL DANGER"
+        elif d.hospital_capacity_remaining < 0.6:
+            hosp_status = "hospital LOW"
+        else:
+            hosp_status = "hospital OK"
+
+        lag_note = " [DATA IS 3 DAYS OLD]" if has_data_lag else ""
         lines.append(
-            f"  District {d.district_id}: "
-            f"infection={d.reported_infection_rate:.2f} [{status}], "
-            f"growth_hint={d.growth_rate_hint:.2f}, "
-            f"hospital={d.hospital_capacity_remaining:.2f}, "
-            f"restricted={'yes' if d.restriction_active else 'no'}, "
-            f"tested_recently={'yes' if d.tested_recently else 'no'}"
+            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})"
         )
 
+    lines += [""]
+
+    if not has_data_lag:
+        # Easy and medium: data is accurate, allocate reduces existing infection
+        lines += [
+            "HOW ACTIONS WORK:",
+            "  - 'allocate': costs 1 resource, REDUCES existing infection AND slows spread",
+            "  - 'restrict': FREE, only slows future spread, does NOT reduce infection",
+            "  - 'test': costs 1 resource, gives accurate data (NOT needed here, data is real-time)",
+            "",
+            "DECISION RULES (follow in order):",
+            "1. If ANY hospital is below 0.3 capacity: 'allocate' on that district IMMEDIATELY.",
+            "2. If resources > 0 and any district is CRITICAL (above 0.4): 'allocate' on the highest.",
+            "3. If resources > 0 and any district is WARNING (0.2-0.4): 'allocate' on the highest.",
+            "4. If resources = 0: 'restrict' on the highest infected district.",
+            "5. NEVER use 'test' — data is already accurate.",
+            "6. NEVER restrict a SAFE district (below 0.2) — you will be penalised.",
+        ]
+    else:
+        # Hard task: data is 3 days old, act on growth_hint
+        lines += [
+            "HOW ACTIONS WORK:",
+            "  - 'allocate': costs 1 resource, reduces infection AND slows spread",
+            "  - 'restrict': FREE, only slows future spread",
+            "  - 'test': costs 1 resource (NOT recommended — data lag is unavoidable)",
+            "",
+            "DECISION RULES (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.",
+            "3. If resources = 0: 'restrict' on the district with highest growth_hint.",
+            "4. NEVER use 'test' — spending resources on data wastes your limited budget.",
+        ]
+
     lines += [
         "",
-        "Available actions:",
-        "  - 'test'     : Get accurate infection data for a district (costs 1 resource)",
-        "  - 'restrict' : Impose movement restriction in a district (free, but penalised if infection is low)",
-        "  - 'allocate' : Deploy medical resources to a district (costs 1 resource)",
-        "",
-        "Strategy hints:",
-        "  - Prioritise districts in DANGER or with high growth_hint",
-        "  - Use 'test' on high growth_hint districts to reveal true infection",
-        "  - Use 'allocate' on the most infected district",
-        "  - Only 'restrict' districts above 0.2 infection rate",
-        "  - If resources = 0, you can only use 'restrict'",
-        "",
-        "Respond with ONLY a JSON object in this exact format:",
-        '{"action_type": "allocate", "district_id": 2}',
+        "Respond with ONLY valid JSON. No explanation. Example:",
+        '{"action_type": "allocate", "district_id": 0}',
         "",
         "Your decision:",
     ]
@@ -80,16 +101,13 @@ def build_prompt(obs: CityObservation) -> str:
     return "\n".join(lines)
 
 
-# ── LLM Call ──────────────────────────────────────────────────────────────────
-
 def call_llm(prompt: str, client: OpenAI) -> str:
-    """Call the LLM and return the raw response string."""
     response = client.chat.completions.create(
         model    = os.environ.get("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct"),
         messages = [
             {
                 "role":    "system",
-                "content": "You are an epidemic response AI. Always respond with valid JSON only. No explanation."
+                "content": "You are an epidemic response AI. Always respond with valid JSON only. No explanation, no markdown."
             },
             {
                 "role":    "user",
@@ -97,27 +115,17 @@ def call_llm(prompt: str, client: OpenAI) -> str:
             }
         ],
         max_tokens  = 50,
-        temperature = 0.2,   # Low temperature for consistent, reliable decisions
+        temperature = 0.1,  # Lower temperature for more consistent decisions
     )
     return (response.choices[0].message.content or "").strip()
 
 
-# ── Response Parser ───────────────────────────────────────────────────────────
-
 def parse_action(response: str, num_districts: int) -> ContainmentAction:
-    """
-    Parse LLM response into a ContainmentAction.
-    Handles common LLM formatting issues defensively.
-    Falls back to a safe default if parsing fails entirely.
-    """
     valid_types = {"test", "restrict", "allocate"}
 
     try:
-        # Strip markdown code fences if present
         cleaned = re.sub(r"```(?:json)?|```", "", response).strip()
-
-        # Extract JSON object if surrounded by other text
-        match = re.search(r"\{.*?\}", cleaned, re.DOTALL)
+        match   = re.search(r"\{.*?\}", cleaned, re.DOTALL)
         if match:
             cleaned = match.group()
 
@@ -125,30 +133,17 @@ def parse_action(response: str, num_districts: int) -> ContainmentAction:
         action_type = str(data.get("action_type", "allocate")).lower().strip()
         district_id = int(data.get("district_id", 0))
 
-        # Validate and clamp
         if action_type not in valid_types:
             action_type = "allocate"
         district_id = max(0, min(district_id, num_districts - 1))
 
-        return ContainmentAction(
-            action_type = action_type,
-            district_id = district_id,
-        )
+        return ContainmentAction(action_type=action_type, district_id=district_id)
 
     except Exception:
-        # Safe fallback — allocate to district 0
         return ContainmentAction(action_type="allocate", district_id=0)
 
 
-# ── Main Policy Function ──────────────────────────────────────────────────────
-
 def get_action(obs: CityObservation, client: OpenAI) -> ContainmentAction:
-    """
-    Main entry point for the policy.
-    Takes an observation, returns a ContainmentAction.
-    Called by evaluator.py on every step.
-    """
-    prompt      = build_prompt(obs)
-    response    = call_llm(prompt, client)
-    action      = parse_action(response, len(obs.districts))
-    return action
+    prompt   = build_prompt(obs)
+    response = call_llm(prompt, client)
+    return parse_action(response, len(obs.districts))

server/constants.py

@@ -2,7 +2,6 @@
 # ─────────────────────────────────────────────────────────────────────────────
 # Single source of truth for all numeric configuration in the environment.
 # Nothing in this file is computed — these are fixed values only.
-# Adjust reward weights here during tuning without touching environment.py.
 # ─────────────────────────────────────────────────────────────────────────────
 
 import sys
@@ -15,56 +14,57 @@ TASK_CONFIG = {
     "easy": {
         "num_districts":    2,
         "max_steps":        10,
-        "resource_pool":    10,     # Resources available per episode
-        "data_lag_days":    0,      # Agent sees real-time infection data
+        "resource_pool":    10,
+        "data_lag_days":    0,
     },
     "medium": {
         "num_districts":    4,
         "max_steps":        15,
-        "resource_pool":    8,      # Tighter budget forces real tradeoffs
+        "resource_pool":    8,
         "data_lag_days":    0,
     },
     "hard": {
         "num_districts":    6,
         "max_steps":        15,
-        "resource_pool":    7,      # Scarce resources + delayed data
-        "data_lag_days":    3,      # Agent sees infection rates from 3 days ago
+        "resource_pool":    7,
+        "data_lag_days":    3,
     },
 }
 
 
 # ── Infection Thresholds ──────────────────────────────────────────────────────
 
-INFECTION_THRESHOLD   = 0.40    # Above this → district is in danger (penalty fires)
-SAFE_THRESHOLD        = 0.20    # Below this → district is contained (bonus fires)
-LOW_THRESHOLD         = 0.20    # Below this → restriction is deemed unnecessary
-HOSPITAL_BREACH_POINT = 0.00    # At or below this → hospital has collapsed
+INFECTION_THRESHOLD   = 0.40
+SAFE_THRESHOLD        = 0.20
+LOW_THRESHOLD         = 0.20
+HOSPITAL_BREACH_POINT = 0.00
 
 
 # ── Spread Mechanics ──────────────────────────────────────────────────────────
 
-SPREAD_RATE_MIN       = 0.04    # Slowest possible true spread rate per day
-SPREAD_RATE_MAX       = 0.15    # Fastest possible true spread rate per day
-GROWTH_HINT_NOISE     = 0.03    # Random noise added to growth_rate_hint (± value)
-TREATMENT_REDUCTION   = 0.06    # Allocating also reduces existing infection
-ALLOCATE_REDUCTION    = 0.10    # How much one 'allocate' reduces spread this step
-RESTRICT_REDUCTION    = 0.05    # How much one 'restrict' reduces spread per step
-SPILLOVER_RATE        = 0.01    # Fraction of infection that spreads to adjacent districts per day
+SPREAD_RATE_MIN       = 0.05    # Slowest possible spread rate per day
+SPREAD_RATE_MAX       = 0.13    # Fastest possible spread rate per day
+GROWTH_HINT_NOISE     = 0.03    # Noise on growth_rate_hint (± value)
+
+TREATMENT_REDUCTION   = 0.08    # Allocating reduces existing infection this step
+ALLOCATE_REDUCTION    = 0.10    # Allocating reduces future spread rate this step
+RESTRICT_REDUCTION    = 0.05    # Restricting reduces spread per step
+SPILLOVER_RATE        = 0.01    # Infection fraction that spills to adjacent districts
 
-RESOURCE_REPLENISH    = 3       # Resource units restored at the start of each new day
+RESOURCE_REPLENISH    = 0       # No replenishment — pool is total budget for episode
 
 
 # ── Reward Weights ────────────────────────────────────────────────────────────
 
-REWARD_INFECTION_PENALTY      = -0.50   # Per district above INFECTION_THRESHOLD each step
-REWARD_HOSPITAL_BREACH        = -1.00   # Per district with breached hospital capacity
-REWARD_EARLY_CONTAINMENT      = +0.50   # Base value; scaled by (1 - step/max_steps)
-REWARD_UNNECESSARY_RESTRICTION = -0.20  # Restricting a district below LOW_THRESHOLD
-REWARD_CORRECT_PRIORITISATION = +0.30   # Allocating to the highest-infected district
+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
 
 
 # ── Episode Terminal Conditions ───────────────────────────────────────────────
 
-# Episode ends early (success) if ALL districts drop below SAFE_THRESHOLD.
-# Episode ends early (failure) if ANY district's hospital capacity hits HOSPITAL_BREACH_POINT.
-# Otherwise episode runs until max_steps is reached.
+# Success: ALL districts below SAFE_THRESHOLD → speed bonus fires
+# Failure: ANY hospital at HOSPITAL_BREACH_POINT → episode ends immediately
+# Natural: max_steps reached

server/grader.py

@@ -125,7 +125,12 @@ def grade_trajectory(
     # Fraction of allocate/test actions that targeted districts above threshold.
     # Rewards directing resources where they're actually needed.
 
-    resource_actions  = [
+    # ── Component 3: Efficiency Score ────────────────────────────────────────
+    # Rewards directing resources to the highest-infected district.
+    # Checks which district had the highest infection at each step,
+    # then rewards targeting it regardless of whether it crossed threshold.
+
+    resource_actions = [
         s for s in trajectory
         if s.action.action_type in {"allocate", "test"}
     ]
@@ -133,12 +138,19 @@ def grade_trajectory(
     if resource_actions:
         correct_actions = 0
         for step in resource_actions:
-            target = step.city_state.districts[step.action.district_id]
-            if target.true_infection_rate > INFECTION_THRESHOLD:
+            districts = step.city_state.districts
+            if not districts:
+                continue
+            # Reward targeting the most infected district at this step
+            highest_id = max(districts, key=lambda d: d.true_infection_rate).district_id
+            if step.action.district_id == highest_id:
+                correct_actions += 1
+            # Also credit actions on districts above threshold (correct triage)
+            elif districts[step.action.district_id].true_infection_rate > INFECTION_THRESHOLD:
                 correct_actions += 1
         efficiency_score = correct_actions / len(resource_actions)
     else:
-        efficiency_score = 0.5  # Neutral if no resource actions taken
+        efficiency_score = 0.5
 
     # ── Component 4: Speed Score ──────────────────────────────────────────────
     # Rewards finishing faster than max_steps.

server/tasks/task_easy.py

@@ -1,14 +1,8 @@
 # server/tasks/task_easy.py
-# ─────────────────────────────────────────────────────────────────────────────
-# Easy task: 2 districts, 1 outbreak, accurate real-time data.
-# Agent should learn to test the infected district, restrict it,
-# and allocate resources. A straightforward strategy scores 0.7–0.9.
-# ─────────────────────────────────────────────────────────────────────────────
-
 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 models import CityState
 from server.utils import generate_districts
@@ -19,15 +13,16 @@ from server.constants import TASK_CONFIG
 class EasyTask(BaseTask):
 
     name          = "easy"
-    num_districts = TASK_CONFIG["easy"]["num_districts"]    # 2
-    max_steps     = TASK_CONFIG["easy"]["max_steps"]        # 10
-    resource_pool = TASK_CONFIG["easy"]["resource_pool"]    # 10
-    data_lag_days = TASK_CONFIG["easy"]["data_lag_days"]    # 0
+    num_districts = TASK_CONFIG["easy"]["num_districts"]
+    max_steps     = TASK_CONFIG["easy"]["max_steps"]
+    resource_pool = TASK_CONFIG["easy"]["resource_pool"]
+    data_lag_days = TASK_CONFIG["easy"]["data_lag_days"]
 
     def build_initial_state(self) -> CityState:
-        # District 0 has a visible outbreak. District 1 is clean.
-        # Agent only needs to identify and respond to one threat.
-        seed_infections = [0.25, 0.05]
+        # D0 has a visible outbreak in WARNING-CRITICAL boundary.
+        # Seed is high enough that one allocation cannot instantly contain it —
+        # agent must sustain 3+ allocations while also managing D1's growth.
+        seed_infections = [0.32, 0.06]
 
         return CityState(
             day                 = 0,
@@ -36,8 +31,8 @@ class EasyTask(BaseTask):
             data_lag_days       = self.data_lag_days,
             max_steps           = self.max_steps,
             districts           = generate_districts(
-                                      num_districts    = self.num_districts,
-                                      seed_infections  = seed_infections,
+                                      num_districts   = self.num_districts,
+                                      seed_infections = seed_infections,
                                   ),
             infection_history   = [],
         )