diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..98439a8697bef9bd99ca21ab2e4ca45fc1ddcde4 Binary files /dev/null and b/.gitignore differ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..481aba5a7154b08481b298ba5b7b3577a304eed8 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,27 @@ +# server/Dockerfile +# ───────────────────────────────────────────────────────────────────────────── +# Builds the Cascade Containment environment server. +# Exposes port 7860 — required for Hugging Face Spaces deployment. +# ───────────────────────────────────────────────────────────────────────────── + +FROM python:3.11-slim + +WORKDIR /app + +RUN apt-get update && apt-get install -y \ + build-essential \ + && rm -rf /var/lib/apt/lists/* + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY models.py . +COPY constants.py . +COPY server/ ./server/ +COPY core/ ./core/ + +ENV PYTHONPATH="/app:/app/server" + +EXPOSE 7860 + +CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"] \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..af27250dd3dd5ff9ed713052b8761a36f5350350 --- /dev/null +++ b/README.md @@ -0,0 +1,53 @@ +--- +title: Cascade Containment +emoji: 🦠 +colorFrom: red +colorTo: blue +sdk: docker +app_port: 7860 +pinned: false +--- + +# Cascade Containment + +An RL benchmark for epidemic containment policy under uncertainty. +A city health authority must allocate limited resources across districts +to contain a spreading outbreak — with delayed data, resource scarcity, +and cascading hospital stress. + +Generalises to wildfire deployment, cyberattack isolation, and misinformation containment. + +## Environment + +- **3 tasks:** Easy (2 districts), Medium (4 districts), Hard (6 districts with 3-day data lag) +- **Action space:** `action_type` (test/restrict/allocate) + `district_id` +- **Learning:** GRPO-style episodic memory with advantage gating + +## Usage + +\```python +from client import CascadeContainmentEnv +from models import ContainmentAction + +with CascadeContainmentEnv(base_url="https://YOUR-SPACE-URL.hf.space").sync() as env: + obs = env.reset(task_name="easy") + result = env.step(ContainmentAction(action_type="allocate", district_id=0)) +\``` + +## Tasks + +| Task | Districts | Steps | Resources | Data Lag | +|------|-----------|-------|-----------|----------| +| easy | 2 | 10 | 10 | None | +| medium | 4 | 15 | 8 | None | +| hard | 6 | 20 | 7 | 3 days | + +## Reward Function + +| Term | Value | Condition | +|------|-------|-----------| +| Infection penalty | -0.50 | Per district above 0.4 threshold | +| Hospital breach | -1.00 | Per breached hospital | +| Early containment | +0.50 | Scaled by time remaining | +| Unnecessary restriction | -0.20 | Restricting below 0.2 threshold | +| Correct prioritisation | +0.30 | Allocating to highest-infected district | \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/__pycache__/client.cpython-313.pyc b/__pycache__/client.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11f5e53d5c4c9dc94425e70f5dc7d475a6f42112 Binary files /dev/null and b/__pycache__/client.cpython-313.pyc differ diff --git a/__pycache__/models.cpython-313.pyc b/__pycache__/models.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ba0be8586c837f3be37bb819b19748f5a671e2b Binary files /dev/null and b/__pycache__/models.cpython-313.pyc differ diff --git a/baseline/__init__.py b/baseline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/baseline/__pycache__/__init__.cpython-313.pyc b/baseline/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ee569f52f3fbf8deaca93ae32d07d5eaa01fdd5 Binary files /dev/null and b/baseline/__pycache__/__init__.cpython-313.pyc differ diff --git a/baseline/__pycache__/evaluator.cpython-313.pyc b/baseline/__pycache__/evaluator.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42112b300306c99cadf5e22a628ea66f698f3848 Binary files /dev/null and b/baseline/__pycache__/evaluator.cpython-313.pyc differ diff --git a/baseline/__pycache__/policy.cpython-313.pyc b/baseline/__pycache__/policy.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b531ab1fc427b2e9b90a224df76ba5d771ee1e9 Binary files /dev/null and b/baseline/__pycache__/policy.cpython-313.pyc differ diff --git a/baseline/evaluator.py b/baseline/evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..c4af3ccae7d0231cfb36fb2f8ddd6fee241ce1c0 --- /dev/null +++ b/baseline/evaluator.py @@ -0,0 +1,197 @@ +# baseline/evaluator.py +# ───────────────────────────────────────────────────────────────────────────── +# GRPO-style evaluation loop for Cascade Containment. +# Imports core components — stays focused on orchestration only. +# ───────────────────────────────────────────────────────────────────────────── + +import os +import sys +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +import time +from typing import List, Tuple +from openai import OpenAI +from typing import Any + +from client import CascadeContainmentEnv +from models import ContainmentAction, CityObservation +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 + +N_ROLLOUTS = 3 + + +# ── 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) + + 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:") + + +# ── Single Rollout ──────────────────────────────────────────────────────────── + +def run_rollout( + 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) + obs = result.observation + done = result.done + total_reward = 0.0 + step = 0 + trajectory = [] + + while not done: + prompt = build_prompt_with_memory(obs, memory) + response = call_llm(prompt, client) + action = parse_action(response, len(obs.districts)) + + result = env.step(action) + next_obs = result.observation + reward = result.reward or 0.0 + done = result.done + total_reward += reward + step += 1 + + 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}" + ) + + obs = next_obs + if done: + break + + return total_reward, step, trajectory + + +# ── GRPO Task Runner ────────────────────────────────────────────────────────── + +def run_task_grpo( + env: Any, + task_name: str, + client: OpenAI, + 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}") + + memory = EpisodicMemory(max_size=20) + rollouts = [] + + for i in range(N_ROLLOUTS): + if verbose: + 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) + score = normalise_score(total_reward, steps) + rollouts.append((total_reward, steps, score)) + + if verbose: + print(f" → Reward: {total_reward:+.4f} | Score: {score:.4f}") + + # GRPO advantage computation + 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 + 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[0])[2] + + if verbose: + print(f"\n Rewards: {[round(r, 4) for r in all_rewards]}") + print(f" Mean: {mean_reward:+.4f}") + print(f" Advantages: {[round(r - mean_reward, 4) for r in all_rewards]}") + print(f" Best score: {best_score:.4f}") + + 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.""" + if verbose: + print("\n" + "="*52) + print(" CASCADE CONTAINMENT — GRPO EVALUATION") + print("="*52) + print(f" Rollouts per task: {N_ROLLOUTS}") + print(f" Learning: Episodic memory + advantage gating") + + client = get_client() + scores = {} + start = time.time() + + with CascadeContainmentEnv(base_url=base_url).sync() as env: + for task_name in ["easy", "medium", "hard"]: + try: + score = run_task_grpo(env, task_name, client, verbose) + scores[task_name] = score + if verbose: + print(f"\n ✓ {task_name.upper()} final score: {score:.4f}") + except Exception as e: + scores[task_name] = 0.0 + if verbose: + print(f" ✗ {task_name.upper()} failed: {e}") + import traceback + traceback.print_exc() + + scores["average"] = round( + sum(v for k, v in scores.items() if k != "average") / 3, + 4 + ) + + elapsed = round(time.time() - start, 1) + + if verbose: + print("\n" + "="*52) + print(" FINAL SCORES") + print("="*52) + print(f" Easy: {scores.get('easy', 0.0):.4f}") + print(f" Medium: {scores.get('medium', 0.0):.4f}") + print(f" Hard: {scores.get('hard', 0.0):.4f}") + print(f" {'─'*32}") + print(f" Average: {scores.get('average', 0.0):.4f}") + print(f" Time: {elapsed}s") + print("="*52 + "\n") + + return scores \ No newline at end of file diff --git a/baseline/policy.py b/baseline/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..9e318faf53061dc67365731a03961c3152d404cc --- /dev/null +++ b/baseline/policy.py @@ -0,0 +1,154 @@ +# 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__), '..')) + +import json +import re +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", ""), + 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. + """ + lines = [ + "You are a public health authority managing an epidemic outbreak.", + "Your goal is to contain infection across all districts before hospitals collapse.", + "", + f"Current situation (Step {obs.current_step}/{obs.max_steps}):", + f"Available resources: {obs.available_resources}", + "", + "District status:", + ] + + 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" + 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'}" + ) + + 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}', + "", + "Your decision:", + ] + + 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." + }, + { + "role": "user", + "content": prompt + } + ], + max_tokens = 50, + temperature = 0.2, # Low temperature for consistent, reliable 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) + if match: + cleaned = match.group() + + data = json.loads(cleaned) + 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, + ) + + 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 \ No newline at end of file diff --git a/baseline/run.py b/baseline/run.py new file mode 100644 index 0000000000000000000000000000000000000000..0240c7739306b935da35f31af454ce3d77d4cf50 --- /dev/null +++ b/baseline/run.py @@ -0,0 +1,32 @@ +# 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 +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 + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/client.py b/client.py new file mode 100644 index 0000000000000000000000000000000000000000..1ab413700607d5a49e699aef27cd4efd526d4bd2 --- /dev/null +++ b/client.py @@ -0,0 +1,66 @@ +# 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 +from models import ContainmentAction, CityObservation + + +class CascadeContainmentEnv(EnvClient[ContainmentAction, CityObservation, State]): + """ + Client for the Cascade Containment OpenEnv environment. + + Async usage: + 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: + 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), + ) + + +# ── Connection test (run directly to verify client works) ───────────────────── + +if __name__ == "__main__": + 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}") + + result = env.step(ContainmentAction(action_type="allocate", district_id=0)) + print(f" Step reward: {result.reward}") + print(f"✓ Client working end-to-end") \ No newline at end of file diff --git a/config/tasks.yaml b/config/tasks.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/core/__init__.py b/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/core/__pycache__/__init__.cpython-313.pyc b/core/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1b25a5ab3f431d10d706b9cce89cc061a7bd86c0 Binary files /dev/null and b/core/__pycache__/__init__.cpython-313.pyc differ diff --git a/core/__pycache__/policy_update.cpython-313.pyc b/core/__pycache__/policy_update.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..241f2d231f712c7effc4b601e97ffa8901fff957 Binary files /dev/null and b/core/__pycache__/policy_update.cpython-313.pyc differ diff --git a/core/__pycache__/reward.cpython-313.pyc b/core/__pycache__/reward.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cdfafda3955f1f545a97ba3eca662c7576ec53c4 Binary files /dev/null and b/core/__pycache__/reward.cpython-313.pyc differ diff --git a/core/__pycache__/trajectory.cpython-313.pyc b/core/__pycache__/trajectory.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5978d353d4c8f24d8947befc130a695496448eda Binary files /dev/null and b/core/__pycache__/trajectory.cpython-313.pyc differ diff --git a/core/policy_update.py b/core/policy_update.py new file mode 100644 index 0000000000000000000000000000000000000000..92f6807e3c527e6bcc5b31a7f4c320170908aa9d --- /dev/null +++ b/core/policy_update.py @@ -0,0 +1,58 @@ +# 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 core.trajectory import EpisodicMemory + + +def compute_advantage( + 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. + """ + if not completed_rewards: + return 0.0 + mean = sum(completed_rewards) / len(completed_rewards) + return round(current_reward - mean, 4) + + +def should_reinforce(advantage: float) -> bool: + """ + Reinforce if advantage >= 0 (at or above mean). + Suppress if below mean. + """ + return advantage >= 0 + + +def update_memory( + memory: EpisodicMemory, + trajectory: List[dict], + 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. + """ + if not should_reinforce(advantage): + return 0 + + stored = 0 + for step_data in trajectory: + if step_data["reward"] > 0: + memory.store(step_data["obs"], step_data["action"], step_data["reward"]) + stored += 1 + + return stored \ No newline at end of file diff --git a/core/reward.py b/core/reward.py new file mode 100644 index 0000000000000000000000000000000000000000..37a586322d1739dc9752b6c16d546b6cfd9e60d4 --- /dev/null +++ b/core/reward.py @@ -0,0 +1,20 @@ +# core/reward.py +# ───────────────────────────────────────────────────────────────────────────── +# Reward computation utilities for the GRPO evaluation loop. +# ───────────────────────────────────────────────────────────────────────────── + +import math + + +def normalise_score(total_reward: float, steps: int) -> float: + """ + Map cumulative reward to [0.0, 1.0] via sigmoid on average reward per step. + Guaranteed to always return a value strictly within the valid range. + + Average reward of 0 → 0.5 + Positive average → above 0.5 + Negative average → below 0.5 + """ + raw = total_reward / max(steps, 1) + score = 1.0 / (1.0 + math.exp(-raw)) + return round(min(1.0, max(0.0, score)), 4) \ No newline at end of file diff --git a/core/trajectory.py b/core/trajectory.py new file mode 100644 index 0000000000000000000000000000000000000000..fe1830de665424f1d9033dd0eef8cd5cc489c7d2 --- /dev/null +++ b/core/trajectory.py @@ -0,0 +1,81 @@ +# 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 +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from typing import List +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" + """ + + def __init__(self, max_size: int = 20): + self.memories: List[dict] = [] + 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: + return + + self.memories.append({ + "infection_profile": [round(d.reported_infection_rate, 2) for d in obs.districts], + "resources": obs.available_resources, + "action_type": action.action_type, + "district_id": action.district_id, + "reward": round(reward, 4), + }) + + # 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: + """ + 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 "" + + current = [round(d.reported_infection_rate, 2) for d in obs.districts] + + def l1_distance(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)) + + ranked = sorted(self.memories, key=l1_distance) + top = ranked[:top_k] + + lines = ["Relevant past decisions (from successful rollouts):"] + 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}" + ) + return "\n".join(lines) + + def clear(self): + """Clear memory between tasks — memories are task-specific.""" + self.memories = [] + + def __len__(self) -> int: + return len(self.memories) \ No newline at end of file diff --git a/inference.py b/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..3ebf9cf0a979360c307030d81c3f99adfd882df3 --- /dev/null +++ b/inference.py @@ -0,0 +1,44 @@ +# inference.py +# ───────────────────────────────────────────────────────────────────────────── +# Root-level entry point for hackathon evaluation. +# Judges run this file to verify reproducible scores across all three tasks. +# +# Required environment variables: +# API_BASE_URL — LLM API endpoint +# MODEL_NAME — Model identifier for inference +# HF_TOKEN — Hugging Face / API key +# ENV_BASE_URL — Running environment server URL (default: localhost:7860) +# +# Usage: +# python inference.py +# +# Runtime must be under 20 minutes on 2vCPU / 8GB RAM. +# ───────────────────────────────────────────────────────────────────────────── + +import os +import sys +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from baseline.run import main + + +if __name__ == "__main__": + scores = main() + + # Machine-readable summary for auto-validator + print("\nSCORES:") + print(f" easy: {scores.get('easy', 0.0):.4f}") + print(f" medium: {scores.get('medium', 0.0):.4f}") + print(f" hard: {scores.get('hard', 0.0):.4f}") + print(f" average: {scores.get('average', 0.0):.4f}") + + # Warn and exit non-zero if evaluation failed entirely + if scores.get("average", 0.0) == 0.0: + print("\nWARNING: All scores are zero. Check:") + print(" 1. Is the environment server running?") + print(f" ENV_BASE_URL = {os.environ.get('ENV_BASE_URL', 'http://localhost:7860')}") + print(" 2. Are API credentials set?") + print(f" API_BASE_URL = {os.environ.get('API_BASE_URL', 'NOT SET')}") + print(f" MODEL_NAME = {os.environ.get('MODEL_NAME', 'NOT SET')}") + print(f" HF_TOKEN = {'SET' if os.environ.get('HF_TOKEN') else 'NOT SET'}") + sys.exit(1) \ No newline at end of file diff --git a/models.py b/models.py new file mode 100644 index 0000000000000000000000000000000000000000..87d9cfdc13278f5f552bdd1e764b1bc3032a0a59 --- /dev/null +++ b/models.py @@ -0,0 +1,70 @@ +from dataclasses import dataclass, field +from typing import List, Optional +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 + + +# ── 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 + hospital_capacity_remaining: float + population_density: float + days_since_tested: int + restriction_active: bool + deployed_resources: int # Resource units currently active here + + +# ── 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. + +@dataclass +class CityState: + day: int = 0 + available_resources: int = 0 + task_name: str = "easy" + data_lag_days: int = 0 + max_steps: int = 10 + districts: List[DistrictTruth] = field(default_factory=list) + 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 + """ + 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. + +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") + max_steps: int = Field(..., description="Total steps allowed this episode") + message: Optional[str] = Field(None, description="Human-readable feedback for debugging") \ No newline at end of file diff --git a/openenv.yaml b/openenv.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c7eeb6f77683bc3fd36556fbbac43810f7ced9a8 --- /dev/null +++ b/openenv.yaml @@ -0,0 +1,120 @@ +# 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: > + An RL benchmark for epidemic containment policy under uncertainty. + A city health authority must allocate limited resources across districts + to contain a spreading outbreak — with delayed data, resource scarcity, + and cascading hospital stress. Generalises to wildfire deployment, + cyberattack isolation, and misinformation containment. + +author: SST-Team +license: MIT + +# ── Environment Entry Point ─────────────────────────────────────────────────── + +server: + module: server.app + app: app + port: 7860 + dockerfile: Dockerfile + +# ── Action Space ────────────────────────────────────────────────────────────── + +action: + type: object + class: ContainmentAction + fields: + action_type: + type: string + description: "One of: 'test', 'restrict', 'allocate'" + enum: [test, restrict, allocate] + district_id: + type: integer + description: "Target district index (0-indexed)" + minimum: 0 + +# ── Observation Space ───────────────────────────────────────────────────────── + +observation: + type: object + class: CityObservation + fields: + districts: + type: array + description: "Per-district state visible to agent" + items: + type: object + fields: + district_id: + type: integer + reported_infection_rate: + type: number + minimum: 0.0 + maximum: 1.0 + growth_rate_hint: + type: number + minimum: 0.0 + maximum: 1.0 + hospital_capacity_remaining: + type: number + minimum: 0.0 + maximum: 1.0 + population_density: + type: number + minimum: 0.0 + maximum: 1.0 + tested_recently: + type: boolean + restriction_active: + type: boolean + available_resources: + type: integer + description: "Resource units remaining this turn" + current_step: + type: integer + max_steps: + type: integer + done: + type: boolean + reward: + type: number + nullable: true + message: + type: string + nullable: true + +# ── Tasks ───────────────────────────────────────────────────────────────────── + +tasks: + - name: easy + description: "2 districts, 1 outbreak, real-time data, generous resources" + max_steps: 10 + num_districts: 2 + + - name: medium + description: "4 districts, 2 simultaneous outbreaks, limited resources" + max_steps: 15 + num_districts: 4 + + - name: hard + description: "6 districts, 3-day data lag, scarce resources" + max_steps: 20 + num_districts: 6 + +# ── Generalisation Note ─────────────────────────────────────────────────────── + +tags: + - reinforcement-learning + - resource-allocation + - sequential-decision-making + - epidemic-containment + - cascade-dynamics + - partial-observability + - openenv \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..ec7b74e9e66b45b09b4f8bd7273da77c9b31106c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[project] +name = "cascade-containment" +version = "1.0.0" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..8e32dcb815830299b0620686ca93407e90f6a6a8 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,7 @@ +# server/requirements.txt +fastapi>=0.104.0 +uvicorn>=0.24.0 +pydantic>=2.0.0 +openenv-core>=0.2.1 +openai>=1.0.0 +python-dotenv>=0.19.0 \ No newline at end of file diff --git a/scripts/test_local.py b/scripts/test_local.py new file mode 100644 index 0000000000000000000000000000000000000000..8c67c7b7bf6308199fa5a75afc22ba5d9af58000 --- /dev/null +++ b/scripts/test_local.py @@ -0,0 +1,131 @@ +# scripts/test_local.py +# Quick sanity check for everything built so far. +# Run this from the project root: python scripts/test_local.py + +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__), '../server')) + +from server.environment import EpidemicContainmentEnv +from models import ContainmentAction +from server.grader import grade_trajectory, grade_task + +def test_grader(task_name: str): + print(f"\n--- Grader test: {task_name} ---") + env = EpidemicContainmentEnv() + obs = env.reset(task_name) + + while not obs.done: + action = ContainmentAction(action_type="allocate", district_id=0) + obs = env.step(action) + + trajectory = env.get_trajectory() + result = grade_trajectory(trajectory, task_name) + + print(f" Final score: {result.final_score:.4f}") + print(f" Containment: {result.containment_score:.4f}") + print(f" Hospital: {result.hospital_score:.4f}") + print(f" Efficiency: {result.efficiency_score:.4f}") + print(f" Speed: {result.speed_score:.4f}") + print(f" Hospital breached: {result.hospital_breached}") + print(f" Districts safe: {result.districts_contained}") + print(f" Steps taken: {result.total_steps}") + assert 0.0 <= result.final_score <= 1.0, "Score out of range!" + print(f"✓ Score in valid range [0.0, 1.0]") + + +def test_task(task_name: str): + print(f"\n{'='*50}") + print(f"Testing task: {task_name.upper()}") + print(f"{'='*50}") + + env = EpidemicContainmentEnv() + + # Test reset() + obs = env.reset(task_name) + print(f"✓ reset() OK") + print(f" Districts: {len(obs.districts)}") + print(f" Resources: {obs.available_resources}") + print(f" Max steps: {obs.max_steps}") + print(f" Message: {obs.message}") + + # Test state() + state = env.state + print(f"✓ state() OK") + print(f" Episode ID: {state.episode_id}") + print(f" Step count: {state.step_count}") + + # Run a few steps with different action types + actions = [ + ContainmentAction(action_type="test", district_id=0), + ContainmentAction(action_type="allocate", district_id=0), + ContainmentAction(action_type="restrict", district_id=1), + ContainmentAction(action_type="allocate", district_id=0), + ContainmentAction(action_type="test", district_id=1), + ] + + total_reward = 0.0 + for i, action in enumerate(actions): + obs = env.step(action) + total_reward += obs.reward or 0.0 + print(f" Step {i+1}: {action.action_type:8} → district {action.district_id} " + f"| reward: {obs.reward:+.4f} | done: {obs.done}") + if obs.done: + print(f" Episode ended early: {obs.message}") + break + + print(f"✓ step() OK — total reward so far: {total_reward:+.4f}") + + # Test invalid action handling + obs = env.reset(task_name) + bad_action = ContainmentAction(action_type="invalid_type", district_id=99) + obs = env.step(bad_action) + print(f"✓ Invalid action handled gracefully: {obs.message}") + + +def run_full_episode(task_name: str): + """Run a complete episode to verify terminal conditions work.""" + print(f"\n--- Full episode: {task_name} ---") + env = EpidemicContainmentEnv() + obs = env.reset(task_name) + + total_reward = 0.0 + step = 0 + + while not obs.done: + # Simple greedy policy: always allocate to district 0 + action = ContainmentAction(action_type="allocate", district_id=0) + obs = env.step(action) + total_reward += obs.reward or 0.0 + step += 1 + + print(f" Ended at step {step}: {obs.message}") + print(f" Total reward: {total_reward:+.4f}") + print(f"✓ Full episode completed cleanly") + + +if __name__ == "__main__": + print("Running Cascade Containment environment tests...\n") + + try: + test_task("easy") + test_task("medium") + test_task("hard") + + test_grader("easy") + test_grader("medium") + test_grader("hard") + + run_full_episode("easy") + run_full_episode("medium") + run_full_episode("hard") + + print(f"\n{'='*50}") + print("✓ ALL TESTS PASSED") + print(f"{'='*50}\n") + + except Exception as e: + print(f"\n✗ TEST FAILED: {e}") + import traceback + traceback.print_exc() \ No newline at end of file diff --git a/scripts/validate.py b/scripts/validate.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/server/__init__.py b/server/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/server/__pycache__/__init__.cpython-313.pyc b/server/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21fda3c1678d5c11fd95fa62f68a0803bee8d504 Binary files /dev/null and b/server/__pycache__/__init__.cpython-313.pyc differ diff --git a/server/__pycache__/app.cpython-313.pyc b/server/__pycache__/app.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46878db40a44c2fe44b7da9551f9dfb35ae9afd4 Binary files /dev/null and b/server/__pycache__/app.cpython-313.pyc differ diff --git a/server/__pycache__/constants.cpython-313.pyc b/server/__pycache__/constants.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d4c7f108db56b2403719e91b2e8405c380225434 Binary files /dev/null and b/server/__pycache__/constants.cpython-313.pyc differ diff --git a/server/__pycache__/environment.cpython-313.pyc b/server/__pycache__/environment.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85eff575cc7cb6f86af3b8780dd272b2bd258edc Binary files /dev/null and b/server/__pycache__/environment.cpython-313.pyc differ diff --git a/server/__pycache__/grader.cpython-313.pyc b/server/__pycache__/grader.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5cb36a1258b79ab5057454caa9e0246b919b5da3 Binary files /dev/null and b/server/__pycache__/grader.cpython-313.pyc differ diff --git a/server/__pycache__/utils.cpython-313.pyc b/server/__pycache__/utils.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c5b107afebabdccfbe1ce62945bda4683855e10 Binary files /dev/null and b/server/__pycache__/utils.cpython-313.pyc differ diff --git a/server/app.py b/server/app.py new file mode 100644 index 0000000000000000000000000000000000000000..586891e1fb71ea7943d570cace18c8b20b4dcbd6 --- /dev/null +++ b/server/app.py @@ -0,0 +1,21 @@ +# server/app.py +# ───────────────────────────────────────────────────────────────────────────── +# FastAPI application entry point for Cascade Containment. +# Uses a factory function so each WebSocket session gets its own isolated +# environment instance — required for concurrent session safety. +# ───────────────────────────────────────────────────────────────────────────── + +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from openenv.core.env_server import create_app +from server.environment import EpidemicContainmentEnv +from models import ContainmentAction, CityObservation + + +app = create_app( + EpidemicContainmentEnv, + ContainmentAction, + CityObservation, +) \ No newline at end of file diff --git a/server/constants.py b/server/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..2f6ae191993cb52c86562a41f63f7be6890e069a --- /dev/null +++ b/server/constants.py @@ -0,0 +1,70 @@ +# constants.py +# ───────────────────────────────────────────────────────────────────────────── +# 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 +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +# ── Task Configuration ──────────────────────────────────────────────────────── + +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 + }, + "medium": { + "num_districts": 4, + "max_steps": 15, + "resource_pool": 8, # Tighter budget forces real tradeoffs + "data_lag_days": 0, + }, + "hard": { + "num_districts": 6, + "max_steps": 20, + "resource_pool": 7, # Scarce resources + delayed data + "data_lag_days": 3, # Agent sees infection rates from 3 days ago + }, +} + + +# ── 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 + + +# ── Spread Mechanics ────────────────────────────────────────────────────────── + +SPREAD_RATE_MIN = 0.05 # Slowest possible true spread rate per day +SPREAD_RATE_MAX = 0.20 # Fastest possible true spread rate per day +GROWTH_HINT_NOISE = 0.03 # Random noise added to growth_rate_hint (± value) + +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.02 # Fraction of infection that spreads to adjacent districts per day + +RESOURCE_REPLENISH = 3 # Resource units restored at the start of each new day + + +# ── 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 + + +# ── 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. \ No newline at end of file diff --git a/server/environment.py b/server/environment.py new file mode 100644 index 0000000000000000000000000000000000000000..3739a00daaa595a6d98645abc9a18b21066edf05 --- /dev/null +++ b/server/environment.py @@ -0,0 +1,316 @@ +# 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__), '../..')) + +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 + +from models import ( + CityState, + CityObservation, + ContainmentAction, +) +from server.constants import ( + TASK_CONFIG, + INFECTION_THRESHOLD, + SAFE_THRESHOLD, + LOW_THRESHOLD, + ALLOCATE_REDUCTION, + RESTRICT_REDUCTION, + RESOURCE_REPLENISH, + REWARD_INFECTION_PENALTY, + REWARD_HOSPITAL_BREACH, + REWARD_EARLY_CONTAINMENT, + REWARD_UNNECESSARY_RESTRICTION, + REWARD_CORRECT_PRIORITISATION, +) +from server.utils import ( + build_observation, + compute_spread, + get_highest_infected_district, + all_districts_contained, + any_hospital_breached, + districts_above_threshold, + snapshot_infection_rates, + generate_episode_id, +) +from server.tasks.registry import get_task + + +class EpidemicContainmentEnv(Environment): + """ + Cascade Containment — an RL environment for epidemic response policy. + + The agent plays a city health authority making sequential resource + allocation decisions under uncertainty and delayed feedback. + + Interface: + reset(task_name) → CityObservation + step(action) → CityObservation + state() → State + """ + + 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 = [] + + # ── 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._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, + ) + + 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 + ) + + # ── 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._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, + )) + + # ── 10. Check terminal conditions ───────────────────────────────────── + done, terminal_message = self._check_terminal() + + # ── 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, + done = done, + ) + + @property + def state(self) -> State: + return self._state + + # ── Private: 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. + """ + 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." + + # 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}." + + # 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 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 + + elif action.action_type == "allocate": + # Deploy one resource unit — reduces spread this step via compute_spread + district.deployed_resources += 1 + 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 ──────────────────────────────────────── + + 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.15, 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 + ) + + # ── 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): + reward += REWARD_INFECTION_PENALTY + + # Term 2: Heavy penalty for hospital capacity breach + for district in self._city.districts: + if district.hospital_capacity_remaining <= 0.0: + 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 \ No newline at end of file diff --git a/server/grader.py b/server/grader.py new file mode 100644 index 0000000000000000000000000000000000000000..5f9b100ac017805c8529ed8e5fa66f70f89e34ee --- /dev/null +++ b/server/grader.py @@ -0,0 +1,197 @@ +# server/grader.py +# ───────────────────────────────────────────────────────────────────────────── +# Deterministic scorer for completed Cascade Containment episodes. +# Called by baseline/evaluator.py after each full episode. +# Always returns a float in [0.0, 1.0]. +# ───────────────────────────────────────────────────────────────────────────── + +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from typing import List, Tuple +from dataclasses import dataclass + +from models import CityState, ContainmentAction +from server.constants import ( + INFECTION_THRESHOLD, + SAFE_THRESHOLD, + HOSPITAL_BREACH_POINT, + TASK_CONFIG, +) + + +# ── Trajectory Record ───────────────────────────────────────────────────────── + +@dataclass +class TrajectoryStep: + """ + A single recorded step in an episode. + Stored by environment.py and passed to the grader after episode ends. + """ + step: int + city_state: CityState # Hidden ground truth at this step + action: ContainmentAction # What the agent did + reward: float # Reward received + done: bool # Was this the final step + + +# ── Grader Score Breakdown ──────────────────────────────────────────────────── + +@dataclass +class GradeResult: + """ + Full scoring breakdown for one episode. + The final_score is what the evaluator reports. + """ + final_score: float # Weighted composite: 0.0 to 1.0 + containment_score: float # How well infection was kept below threshold + hospital_score: float # How well hospital capacity was preserved + efficiency_score: float # How well resources were directed + speed_score: float # How quickly the episode was resolved + hospital_breached: bool # Whether any hospital collapse occurred + districts_contained: int # How many districts ended below safe threshold + total_steps: int # Steps taken before episode ended + + +# ── Main Grader ─────────────────────────────────────────────────────────────── + +def grade_trajectory( + trajectory: List[TrajectoryStep], + task_name: str, +) -> GradeResult: + """ + Score a completed episode trajectory. + + Args: + trajectory: Ordered list of TrajectoryStep from one full episode. + task_name: "easy", "medium", or "hard" — affects scoring strictness. + + Returns: + GradeResult with final_score in [0.0, 1.0] and full breakdown. + """ + if not trajectory: + return GradeResult( + final_score = 0.0, + containment_score = 0.0, + hospital_score = 0.0, + efficiency_score = 0.0, + speed_score = 0.0, + hospital_breached = False, + districts_contained = 0, + total_steps = 0, + ) + + config = TASK_CONFIG[task_name] + num_districts = config["num_districts"] + max_steps = config["max_steps"] + total_steps = len(trajectory) + + # ── Component 1: Containment Score ─────────────────────────────────────── + # Fraction of district-days that stayed below infection threshold. + # Perfect agent = 1.0 (no district ever exceeded threshold). + + total_district_days = total_steps * num_districts + safe_district_days = 0 + + for step in trajectory: + for district in step.city_state.districts: + if district.true_infection_rate <= INFECTION_THRESHOLD: + safe_district_days += 1 + + containment_score = safe_district_days / total_district_days + + # ── Component 2: Hospital Score ─────────────────────────────────────────── + # Measures how well hospital capacity was preserved across the episode. + # Any breach = heavy penalty. Near-breach is also penalised proportionally. + + 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 + + avg_capacity = total_capacity_preserved / total_district_days + hospital_score = avg_capacity * (0.3 if hospital_breached else 1.0) + hospital_score = round(min(1.0, max(0.0, hospital_score)), 4) + + # ── Component 3: Efficiency Score ──────────────────────────────────────── + # Fraction of allocate/test actions that targeted districts above threshold. + # Rewards directing resources where they're actually needed. + + resource_actions = [ + s for s in trajectory + if s.action.action_type in {"allocate", "test"} + ] + + 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: + correct_actions += 1 + efficiency_score = correct_actions / len(resource_actions) + else: + efficiency_score = 0.5 # Neutral if no resource actions taken + + # ── Component 4: Speed Score ────────────────────────────────────────────── + # Rewards finishing faster than max_steps. + # If episode ran to max_steps, speed_score = 0.0. + # If contained in half the steps, speed_score = 0.5. Etc. + + last_step = trajectory[-1] + if last_step.done and not hospital_breached: + speed_score = round(1.0 - (total_steps / max_steps), 4) + speed_score = max(0.0, speed_score) + else: + speed_score = 0.0 # No speed bonus for failed or incomplete episodes + + # ── Final Weighted Score ────────────────────────────────────────────────── + # Weights reflect judging priorities: + # containment = primary signal + # hospital = safety constraint + # efficiency = quality differentiator + # speed = tiebreaker + + final_score = ( + containment_score * 0.45 + + hospital_score * 0.30 + + efficiency_score * 0.15 + + speed_score * 0.10 + ) + final_score = round(min(1.0, max(0.0, final_score)), 4) + + # ── Final district count ────────────────────────────────────────────────── + final_step = trajectory[-1] + districts_contained = sum( + 1 for d in final_step.city_state.districts + if d.true_infection_rate < SAFE_THRESHOLD + ) + + return GradeResult( + final_score = final_score, + containment_score = round(containment_score, 4), + hospital_score = hospital_score, + efficiency_score = round(efficiency_score, 4), + speed_score = speed_score, + hospital_breached = hospital_breached, + districts_contained = districts_contained, + total_steps = total_steps, + ) + + +# ── Convenience: Grade a Single Score to 0.0–1.0 ───────────────────────────── + +def grade_task( + trajectory: List[TrajectoryStep], + task_name: str, +) -> float: + """ + Thin wrapper that returns just the final_score float. + Used by baseline/evaluator.py for clean score reporting. + """ + result = grade_trajectory(trajectory, task_name) + return result.final_score \ No newline at end of file diff --git a/server/tasks/__init__.py b/server/tasks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/server/tasks/__pycache__/__init__.cpython-313.pyc b/server/tasks/__pycache__/__init__.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..911fa947867d3007c2b378b8fc7497fd0bf4de53 Binary files /dev/null and b/server/tasks/__pycache__/__init__.cpython-313.pyc differ diff --git a/server/tasks/__pycache__/base.cpython-313.pyc b/server/tasks/__pycache__/base.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8121d7fc7f9493761ef3430790982e118392cbe Binary files /dev/null and b/server/tasks/__pycache__/base.cpython-313.pyc differ diff --git a/server/tasks/__pycache__/registry.cpython-313.pyc b/server/tasks/__pycache__/registry.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9cb7a801b81e938e7ffefdbbafd7b81f80567735 Binary files /dev/null and b/server/tasks/__pycache__/registry.cpython-313.pyc differ diff --git a/server/tasks/__pycache__/task_easy.cpython-313.pyc b/server/tasks/__pycache__/task_easy.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f27a25cec80bb0e60f4c5a9ba312c2d672a43528 Binary files /dev/null and b/server/tasks/__pycache__/task_easy.cpython-313.pyc differ diff --git a/server/tasks/__pycache__/task_hard.cpython-313.pyc b/server/tasks/__pycache__/task_hard.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c7f20db67f8b523094beaf5be26a9c73e38e894 Binary files /dev/null and b/server/tasks/__pycache__/task_hard.cpython-313.pyc differ diff --git a/server/tasks/__pycache__/task_medium.cpython-313.pyc b/server/tasks/__pycache__/task_medium.cpython-313.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f8f99f1e2d43df89b457c7ac653e62fbc9ab359f Binary files /dev/null and b/server/tasks/__pycache__/task_medium.cpython-313.pyc differ diff --git a/server/tasks/base.py b/server/tasks/base.py new file mode 100644 index 0000000000000000000000000000000000000000..6f2bfa4af4411a4cb7c84059ee61f12b4de9794e --- /dev/null +++ b/server/tasks/base.py @@ -0,0 +1,33 @@ +# 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__), '../..')) + +from abc import ABC, abstractmethod +from models import CityState + + +class BaseTask(ABC): + + # These must be defined by every subclass + name: str + num_districts: int + max_steps: int + resource_pool: int + data_lag_days: int + + @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(). + """ + ... + + def __repr__(self) -> str: + return f"Task(name={self.name}, districts={self.num_districts}, steps={self.max_steps})" \ No newline at end of file diff --git a/server/tasks/registry.py b/server/tasks/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..9d268c81410f187878a48e6f12d62ff355206e7d --- /dev/null +++ b/server/tasks/registry.py @@ -0,0 +1,38 @@ +# 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 + +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, + "medium": MediumTask, + "hard": HardTask, +} + + +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]() \ No newline at end of file diff --git a/server/tasks/task_easy.py b/server/tasks/task_easy.py new file mode 100644 index 0000000000000000000000000000000000000000..03b0ebf6c5b1ac2aab6a62d0ec768a7aee3399ac --- /dev/null +++ b/server/tasks/task_easy.py @@ -0,0 +1,43 @@ +# 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 + +from models import CityState +from server.utils import generate_districts +from server.tasks.base import BaseTask +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 + + 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] + + return CityState( + day = 0, + available_resources = self.resource_pool, + task_name = self.name, + data_lag_days = self.data_lag_days, + max_steps = self.max_steps, + districts = generate_districts( + num_districts = self.num_districts, + seed_infections = seed_infections, + ), + infection_history = [], + ) \ No newline at end of file diff --git a/server/tasks/task_hard.py b/server/tasks/task_hard.py new file mode 100644 index 0000000000000000000000000000000000000000..cf5a92186fae704cd5a7e9904d86dac5e47cc06d --- /dev/null +++ b/server/tasks/task_hard.py @@ -0,0 +1,54 @@ +# server/tasks/task_hard.py +# ───────────────────────────────────────────────────────────────────────────── +# Hard task: 6 districts, 3-day data lag, scarce resources. +# All districts start with small but growing infections. +# Agent must learn to read growth_rate_hint signals and act proactively +# on districts that look manageable today but will be critical in 3 days. +# ───────────────────────────────────────────────────────────────────────────── + +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 + +from models import CityState +from server.utils import generate_districts +from server.tasks.base import BaseTask +from server.constants import TASK_CONFIG + + +class HardTask(BaseTask): + + name = "hard" + num_districts = TASK_CONFIG["hard"]["num_districts"] # 6 + max_steps = TASK_CONFIG["hard"]["max_steps"] # 20 + resource_pool = TASK_CONFIG["hard"]["resource_pool"] # 7 + data_lag_days = TASK_CONFIG["hard"]["data_lag_days"] # 3 + + def build_initial_state(self) -> CityState: + # All districts start with small infections that grow at different rates. + # The 3-day lag means the agent won't see today's true rates until day 3. + # Districts with high true_spread_rate will accelerate invisibly. + seed_infections = [0.10, 0.08, 0.12, 0.07, 0.15, 0.09] + + # Pre-populate infection_history with 3 days of identical + # starting values so the lag mechanic works from step 1. + initial_rates = seed_infections[:] + infection_history = [ + initial_rates[:], # day -3 (what agent sees on step 1) + initial_rates[:], # day -2 + initial_rates[:], # day -1 + ] + + return CityState( + day = 0, + available_resources = self.resource_pool, + task_name = self.name, + data_lag_days = self.data_lag_days, + max_steps = self.max_steps, + districts = generate_districts( + num_districts = self.num_districts, + seed_infections = seed_infections, + ), + infection_history = infection_history, + ) \ No newline at end of file diff --git a/server/tasks/task_medium.py b/server/tasks/task_medium.py new file mode 100644 index 0000000000000000000000000000000000000000..3b86a6fd861983e02860c056960f1cd5ad4d9c9d --- /dev/null +++ b/server/tasks/task_medium.py @@ -0,0 +1,44 @@ +# server/tasks/task_medium.py +# ───────────────────────────────────────────────────────────────────────────── +# Medium task: 4 districts, 2 simultaneous outbreaks, tighter resource pool. +# Agent must prioritise between competing threats — it cannot fully +# address both outbreaks simultaneously and must learn to triage. +# ───────────────────────────────────────────────────────────────────────────── + +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 + +from models import CityState +from server.utils import generate_districts +from server.tasks.base import BaseTask +from server.constants import TASK_CONFIG + + +class MediumTask(BaseTask): + + name = "medium" + num_districts = TASK_CONFIG["medium"]["num_districts"] # 4 + max_steps = TASK_CONFIG["medium"]["max_steps"] # 15 + resource_pool = TASK_CONFIG["medium"]["resource_pool"] # 8 + data_lag_days = TASK_CONFIG["medium"]["data_lag_days"] # 0 + + def build_initial_state(self) -> CityState: + # Districts 0 and 2 are seeded with outbreaks (non-adjacent). + # Districts 1 and 3 are clean but will receive spillover. + # Agent must choose which outbreak to tackle first. + seed_infections = [0.30, 0.05, 0.28, 0.05] + + return CityState( + day = 0, + available_resources = self.resource_pool, + task_name = self.name, + data_lag_days = self.data_lag_days, + max_steps = self.max_steps, + districts = generate_districts( + num_districts = self.num_districts, + seed_infections = seed_infections, + ), + infection_history = [], + ) \ No newline at end of file diff --git a/server/utils.py b/server/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..67a33bfb8d6753038958f772576bff42b01f8134 --- /dev/null +++ b/server/utils.py @@ -0,0 +1,216 @@ +# utils.py +# ───────────────────────────────────────────────────────────────────────────── +# Helper functions used by environment.py, grader.py, and task files. +# No game logic lives here — only pure utility functions. +# ───────────────────────────────────────────────────────────────────────────── + +import random +import uuid +from typing import List, Optional +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..')) + +from models import ( + DistrictObservation, + DistrictTruth, + CityObservation, + CityState, +) +from server.constants import ( + SPREAD_RATE_MIN, + SPREAD_RATE_MAX, + GROWTH_HINT_NOISE, + INFECTION_THRESHOLD, + SAFE_THRESHOLD, + HOSPITAL_BREACH_POINT, +) + + +# ── City Generation ─────────────────────────────────────────────────────────── + +def generate_districts( + num_districts: int, + seed_infections: List[float], +) -> List[DistrictTruth]: + """ + Create a fresh list of DistrictTruth objects for a new episode. + + seed_infections is a list of starting infection rates, one per district. + The task files control this — easy seeds one district, medium seeds two, + hard seeds multiple with small values that grow over time. + + Population densities are assigned so they sum to exactly 1.0 across all + districts, reflecting real cities where denser districts carry more risk. + """ + assert len(seed_infections) == num_districts, ( + f"Expected {num_districts} seed values, got {len(seed_infections)}" + ) + + # Generate random population densities that sum to 1.0 + raw_densities = [random.uniform(0.5, 1.5) for _ in range(num_districts)] + total = sum(raw_densities) + densities = [round(d / total, 4) for d in raw_densities] + + districts = [] + for i in range(num_districts): + districts.append(DistrictTruth( + district_id = i, + true_infection_rate = seed_infections[i], + true_spread_rate = round(random.uniform(SPREAD_RATE_MIN, SPREAD_RATE_MAX), 4), + hospital_capacity_remaining = 1.0, + population_density = densities[i], + days_since_tested = 99, # Large value → not recently tested + restriction_active = False, + deployed_resources = 0, + )) + + return districts + + +def generate_episode_id() -> str: + """Return a unique identifier for a new episode.""" + return str(uuid.uuid4())[:8] + + +# ── Observation Builder ─────────────────────────────────────────────────────── + +def build_observation( + state: CityState, + step_count: int, + reward: Optional[float] = None, + message: Optional[str] = None, + done: bool = False, +) -> CityObservation: + """ + Convert the hidden CityState into the CityObservation the agent receives. + + This is where partial observability is enforced. When data_lag_days > 0, + infection rates are pulled from infection_history rather than the current + true values. Everything else (hospital capacity, growth hint, flags) is + always reported in real time. + """ + district_observations = [] + + for i, district in enumerate(state.districts): + + # Apply data lag for hard task + if state.data_lag_days > 0 and len(state.infection_history) >= state.data_lag_days: + reported_rate = state.infection_history[-state.data_lag_days][i] + else: + reported_rate = district.true_infection_rate + + # Add noise to spread rate hint — agent sees a signal, not the truth + noise = random.uniform(-GROWTH_HINT_NOISE, GROWTH_HINT_NOISE) + growth_hint = round(max(0.0, min(1.0, district.true_spread_rate + noise)), 4) + + district_observations.append(DistrictObservation( + district_id = district.district_id, + reported_infection_rate = round(reported_rate, 4), + growth_rate_hint = growth_hint, + hospital_capacity_remaining = round(district.hospital_capacity_remaining, 4), + population_density = district.population_density, + tested_recently = district.days_since_tested <= 2, + restriction_active = district.restriction_active, + )) + + return CityObservation( + districts = district_observations, + available_resources = state.available_resources, + current_step = step_count, + max_steps = state.max_steps, + done = done, + reward = reward, + message = message, + ) + + +# ── Spread Computation ──────────────────────────────────────────────────────── + +def compute_spread(districts: List[DistrictTruth]) -> List[float]: + """ + Calculate new infection rates for all districts after one day passes. + + Each district's infection grows by its spread rate, reduced by any + active restriction or deployed resources. Adjacent districts (by index) + receive a small spillover from their neighbours, simulating geographic + spread without requiring a full spatial model. + + Returns a list of new infection rates (not yet applied to state). + """ + from server.constants import ALLOCATE_REDUCTION, RESTRICT_REDUCTION, SPILLOVER_RATE + + n = len(districts) + new_rates = [] + + for i, district in enumerate(districts): + + # Base growth this day + effective_spread = district.true_spread_rate + + # Reduce spread based on active interventions + if district.restriction_active: + effective_spread = max(0.0, effective_spread - RESTRICT_REDUCTION) + + if district.deployed_resources > 0: + effective_spread = max( + 0.0, + effective_spread - (ALLOCATE_REDUCTION * district.deployed_resources) + ) + + # Grow infection by effective spread rate + new_rate = district.true_infection_rate + effective_spread + + # Add spillover from adjacent districts (wrap-around neighbours) + left_neighbour = districts[(i - 1) % n] + right_neighbour = districts[(i + 1) % n] + + new_rate += left_neighbour.true_infection_rate * SPILLOVER_RATE + new_rate += right_neighbour.true_infection_rate * SPILLOVER_RATE + + # Clamp to valid range + new_rates.append(round(min(1.0, max(0.0, new_rate)), 4)) + + return new_rates + + +# ── Query Helpers ───────────────────────────────────────────────────────────── + +def get_highest_infected_district(districts: List[DistrictTruth]) -> int: + """ + Return the district_id of the district with the highest true infection rate. + Used by the reward function to determine correct prioritisation. + """ + return max(districts, key=lambda d: d.true_infection_rate).district_id + + +def all_districts_contained(districts: List[DistrictTruth]) -> bool: + """ + Return True if every district's true infection rate is below SAFE_THRESHOLD. + This triggers early episode termination with a success condition. + """ + return all(d.true_infection_rate < SAFE_THRESHOLD for d in districts) + + +def any_hospital_breached(districts: List[DistrictTruth]) -> bool: + """ + Return True if any district's hospital capacity has hit zero. + This triggers early episode termination with a failure condition. + """ + return any(d.hospital_capacity_remaining <= HOSPITAL_BREACH_POINT for d in districts) + + +def districts_above_threshold(districts: List[DistrictTruth]) -> List[DistrictTruth]: + """ + Return all districts currently above INFECTION_THRESHOLD. + Used by the reward function to compute per-step infection penalties. + """ + return [d for d in districts if d.true_infection_rate > INFECTION_THRESHOLD] + + +def snapshot_infection_rates(districts: List[DistrictTruth]) -> List[float]: + """ + Return current true infection rates as a plain list, ordered by district_id. + Used by environment.py to append to infection_history each step. + """ + return [d.true_infection_rate for d in sorted(districts, key=lambda d: d.district_id)] \ No newline at end of file diff --git a/structure.txt b/structure.txt new file mode 100644 index 0000000000000000000000000000000000000000..b99d597a41611851911c29062d1971365e8484b1 Binary files /dev/null and b/structure.txt differ