Sync from GitHub

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+episodes.json filter=lfs diff=lfs merge=lfs -text
+sample_episodes.json filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,210 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+#Claude
+.claude/*

Dockerfile.notebook

@@ -0,0 +1,26 @@
+FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
+
+WORKDIR /app
+
+# Install Jupyter + training deps
+RUN pip install --no-cache-dir \
+    jupyter \
+    unsloth \
+    trl \
+    transformers \
+    datasets \
+    accelerate \
+    peft \
+    huggingface_hub
+
+# Copy everything needed for training
+COPY episodes.json .
+COPY generate_episodes.py .
+COPY models.py .
+COPY messages.py .
+COPY drift_events.py .
+COPY notebooks/ ./notebooks/
+
+EXPOSE 8888
+
+CMD ["jupyter", "notebook", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token=''"]

README.md

@@ -1,15 +1,166 @@
----
-title: CrisisInbox
-emoji: 🚨
-colorFrom: red
-colorTo: yellow
-sdk: docker
-app_port: 8000
----
+# CrisisInbox
 
-# CrisisInbox Environment
+A reinforcement learning environment built on [OpenEnv 0.2.1](https://github.com/OpenEnvs/OpenEnv) for training language models to manage personal task overload during natural disasters.
 
-An RL environment for training LLMs to manage personal task overload during natural disasters.
-Built on OpenEnv 0.2.1.
+**Problem Statement 3.2** (Personalized Tasks) + **Patronus AI Sub-Theme** (Schema Drift)
 
-76 messages | 19 senders | 5 schema drift events | 48-hour timeline
+**HF Space:** [eptan-crisis-inbox.hf.space](https://eptan-crisis-inbox.hf.space)
+
+## The Problem
+
+When disaster strikes, your phone explodes. Evacuation orders, panicked family texts, insurance deadlines, your boss demanding slides, your sister begging you to pick up her kids from school, your dad's heart medication left behind. Everything is urgent. Policies change mid-crisis. There are no clean answers — only tradeoffs.
+
+CrisisInbox trains an agent to make those tradeoffs well.
+
+## How It Works
+
+The agent manages a **48-hour post-disaster inbox** as a working parent in Sacramento during a hurricane. Messages arrive in real time from 19 senders across 6 channels. The agent must triage, respond, and adapt — while the rules keep changing underneath.
+
+### Three Layers of Difficulty
+
+**1. Cognitive Overload** — 76 messages arrive over 48 hours. Reading costs time (6 min). Responding costs more (15 min). The agent can't handle everything — it must prioritize and let some things slide.
+
+**2. Conflicting Obligations** — Your boss says come in. HR says take emergency leave. Your sister needs you to watch her kids. Mom wants you to drive to Tahoe. The evacuation shelter is full. There's no right answer, only better tradeoffs.
+
+**3. Schema Drift** — Mid-episode, the rules change:
+- Insurance deadline shortened from 72h to 48h
+- Evacuation zone expands to include your workplace
+- Employer switches from "use PTO" to "5 days paid emergency leave"
+- Airline extends free rebooking from 48h to 7 days
+- FEMA adds new documentation requirements
+
+Each episode randomly fires 3 of 5 drift events. The agent must detect changes and reprioritize.
+
+### Sender Profiles
+
+| Sender | Messages | Tone | Stakes |
+|--------|----------|------|--------|
+| Mom | 8 | Panicked, crying voicemails | Dad's heart medication, family safety |
+| Sister | 7 | Desperate, grateful | Kids stranded at school, childcare |
+| Emma (niece, 7) | 3 | Kid texting style | Scared in the dark, rainbow drawings |
+| Boss (Greg) | 5 | Passive-aggressive, then softens | Career pressure vs. emergency |
+| Neighbor Dave | 8 | Casual bro, community | Cat rescue, looting, cleanup |
+| FEMA / NWS | 10 | Formal, information-dense | Evacuation orders, shelter locations |
+| State Farm | 5 | Corporate | Claim deadlines, documentation |
+| Delta Airlines | 3 | Automated | Flight rebooking |
+| Oakwood Elementary | 3 | School admin | Closures, virtual learning |
+| HR Department | 3 | Policy updates | Leave policies |
+| + 9 others | 21 | Various | Pharmacy, landlord, utilities, etc. |
+
+### MCP Tools (Agent Actions)
+
+| Tool | Time Cost | Description |
+|------|-----------|-------------|
+| `get_inbox` | 0h | View all arrived messages with metadata |
+| `read_message` | 0.1h | Read full content of a specific message |
+| `respond_to_message` | 0.25h | Take action on a message (earns reward) |
+| `get_status` | 0h | View time, score, deadlines, drift events |
+| `advance_time` | 0.5-4h | Skip forward (new messages may arrive) |
+
+### Reward Function
+
+| Signal | Weight | Mechanic |
+|--------|--------|----------|
+| **Urgency base** | 1-10 | Critical=10, High=5, Medium=3, Low=1 |
+| **Deadline timing** | x0.25 to x1.5 | Earlier response = bigger bonus; late = 75% penalty |
+| **Drift adaptation** | x1.5 | Bonus for handling drift-flagged messages |
+| **Stale info penalty** | x0.5 | Penalty for acting on superseded information |
+| **Response quality** | x0.5 | Penalty for very short/empty responses |
+
+An optimal agent responding to the critical evacuation alert at hour 0 earns 15 points. The same response at hour 10 (after the 6h deadline) earns only 2.5.
+
+### Episode Variation
+
+Each episode has:
+- 3 of 5 drift events randomly selected (seed-controlled)
+- +/-15% jitter on message arrival times
+- +/-10% jitter on deadlines
+- Dependency chains that gate actions (e.g., must handle sister's request before school pickup confirmation)
+
+## Quick Start (Hosted)
+
+```python
+from crisis_inbox import CrisisInboxEnv
+
+with CrisisInboxEnv(base_url="https://eptan-crisis-inbox.hf.space") as env:
+    env.reset()
+    inbox = env.call_tool("get_inbox")
+    print(inbox)
+```
+
+## Local Development
+
+```bash
+git clone https://github.com/eptan/crisis-inbox.git
+cd crisis-inbox
+
+python3.12 -m venv .venv
+source .venv/bin/activate
+pip install --upgrade pip
+pip install -e .
+
+# Run server
+uvicorn server.app:app --host 0.0.0.0 --port 8000
+```
+
+Test with the client:
+
+```python
+from crisis_inbox import CrisisInboxEnv
+
+with CrisisInboxEnv(base_url="http://localhost:8000") as env:
+    env.reset()
+
+    # View the inbox
+    inbox = env.call_tool("get_inbox")
+    print(inbox)
+
+    # Read a message
+    msg = env.call_tool("read_message", message_id="msg_001")
+    print(msg)
+
+    # Respond to it
+    result = env.call_tool("respond_to_message",
+        message_id="msg_001",
+        response="Evacuating to Lincoln High School immediately with documents and medication.")
+    print(result)
+
+    # Advance time and check status
+    env.call_tool("advance_time", hours=4.0)
+    status = env.call_tool("get_status")
+    print(status)
+```
+
+## Repository Structure
+
+```
+crisis-inbox/
+├── models.py                           # Message data model (Channel, Urgency, Message)
+├── messages.py                         # 76 pre-written messages across 48h timeline
+├── drift_events.py                     # 5 schema drift events (3 fire per episode)
+├── client.py                           # MCPToolClient subclass
+├── __init__.py                         # Package exports
+├── server/
+│   ├── crisis_inbox_environment.py     # MCPEnvironment with timeline engine
+│   ├── app.py                          # FastAPI app with MCPAction workaround
+│   └── Dockerfile                      # HF Spaces deployment
+├── notebooks/
+│   └── crisisinbox_grpo.ipynb          # GRPO training notebook
+├── episodes.json                       # Pre-generated training episodes
+├── generate_episodes.py                # Episode generator script
+├── pyproject.toml                      # Package config
+├── openenv.yaml                        # OpenEnv environment spec
+├── requirements.txt                    # Docker build dependencies
+└── ROADMAP.md                          # Hackathon timeline and progress
+```
+
+## Stack
+
+- **Environment:** OpenEnv 0.2.1 (MCPEnvironment + FastMCP)
+- **Deployment:** HF Spaces (Docker)
+- **Training:** Unsloth GRPO via Google Colab
+- **Model:** Qwen2.5-0.5B-Instruct
+
+## Team
+
+Built at the OpenEnv Hackathon @ Shack15, SF — March 7-8, 2026

ROADMAP.md

@@ -0,0 +1,252 @@
+# CrisisInbox Hackathon Roadmap
+
+**Event:** OpenEnv Hackathon @ Shack15, SF
+**Dates:** Saturday March 7 – Sunday March 8, 2026
+**Team Size:** 2
+**Problem Statement:** 3.2 (Personalized Tasks) + Patronus AI Sub-Theme (Schema Drift)
+
+---
+
+## PHASE 1: Foundation (11:30 AM – 2:00 PM Saturday)
+
+### Person A — Environment Setup
+
+- [x] Read OpenEnv 0.2.1 docs thoroughly — understand environment structure, observation/action spaces, step function
+- [x] Get a minimal "hello world" OpenEnv environment running locally
+- [x] Define the core data model:
+  - [x] `Message` object: sender, channel (6 types), content, urgency (critical/high/medium/low), deadline, dependencies, drift_flag, supersedes
+  - [x] Task state tracked via environment internals (`_handled`, `_visible_messages`, deadline expiry) — no separate Task object needed since messages and tasks are 1:1
+  - [x] World state tracked via environment instance vars (`_current_hour`, `_all_messages`, `_visible_messages`, `_handled`, `_score`, `_fired_drifts`, `_superseded`) — exposed to agent through `get_status` tool
+- [x] Deploy bare-bones environment to HF Spaces — confirmed at eptan-crisis-inbox.hf.space
+- [ ] Push initial scaffold to GitHub
+
+### Person B — Training Pipeline Setup
+
+- [x] Open Unsloth GRPO Colab notebook and run with a toy example end-to-end
+- [x] Confirm training loop works: environment → agent action → reward → update
+- [x] Define the action space — simplified to free-form: model outputs `respond_to_message(msg_id, "response")`, parsed via regex
+- [x] Define observation space: full inbox snapshot as text prompt with urgency grouping, deadline warnings, drift flags, stale markers
+- [x] Implement reward function (integrated in notebook `score_action()`):
+  - [x] Urgency base: critical=10, high=5, medium=3, low=1
+  - [x] Deadline timing: early bonus (up to +50%), late penalty (-75%)
+  - [x] Schema drift adaptation: +50% for handling drift-flagged messages
+  - [x] Stale info penalty: -50% for acting on superseded messages
+  - [x] Response quality: -50% for short/empty responses
+  - [x] Priority penalty: -70% for choosing low-urgency when critical messages exist
+- [ ] Push training scaffold to GitHub
+
+### Checkpoint — 2:00 PM
+
+- [x] ✅ OpenEnv minimal environment runs locally
+- [x] ✅ HF Spaces deployment pipeline confirmed
+- [x] ✅ Unsloth GRPO training loop configured in notebook (Qwen2.5-0.5B + LoRA + GRPO)
+- [x] ✅ Data model and reward function defined
+- [ ] 🚨 If OpenEnv is blocking: simplify to gym-style env, wrap in OpenEnv later
+- [ ] 🚨 If Unsloth is blocking: fall back to HF TRL directly
+
+---
+
+## PHASE 2: Core Build (2:00 PM – 6:00 PM Saturday)
+
+### Person A — Environment Logic
+
+- [x] Build message generation system:
+  - [x] Create 18 sender profiles (Mom, Boss, Sister, Neighbor Dave, FEMA, NWS, State Farm, Delta Airlines, Oakwood Elementary, HR, Landlord, Coworker Sarah, PG&E, Red Cross, Bank of America, CVS Pharmacy, Sacramento County, Comcast)
+  - [x] Write 76 messages across 48-hour simulated timeline
+  - [x] Each message has: sender, channel, timestamp, urgency, deadline, content, dependencies, drift_flag, supersedes
+  - [x] Organized into waves: initial crisis (hour 0-2), escalation (2-6), post-evacuation chaos (6-12), conflicting demands (12-20), drift events (20-30), recovery (26-36), ongoing management (32-40), final stretch (44-48)
+- [x] Implement episode timeline engine:
+  - [x] Messages arrive at scheduled sim-hours
+  - [x] Time costs per action: reading = 0.1h, responding = 0.25h, advance_time tool = 0.5-4.0h
+  - [x] Dependencies gate actions (must handle prerequisites first)
+  - [x] Episode ends at hour 48
+- [x] Build 5 schema drift events:
+  - [x] Hour 20: Insurance deadline shortened from 72h to 48h
+  - [x] Hour 21: Evacuation zone expanded to include Zone B (workplace)
+  - [x] Hour 22.5: Employer emergency leave expanded (PTO → 5 days paid leave)
+  - [x] Hour 24.5: Airline extends free rebooking window from 48h to 7 days
+  - [x] Hour 34: FEMA adds new documentation requirements
+- [x] Randomization: each episode triggers 3 of 5 drift events (seed-controlled)
+- [x] Add episode variation: +/-15% jitter on arrival times, +/-10% on deadlines (seed-controlled)
+- [x] Integrate with OpenEnv API: `reset()`, `step()` implemented; 5 MCP tools (get_inbox, read_message, respond_to_message, get_status, advance_time)
+
+### Person B — Reward & Training Integration
+
+- [x] Reward function implemented in two places:
+  - [x] Environment-side: `_calculate_reward()` in `crisis_inbox_environment.py` (used during live episodes)
+  - [x] Training-side: `score_action()` in notebook (parses model output, scores against inbox state)
+- [x] Episode generator (`generate_episodes.py`) produces offline training data:
+  - [x] 50 episodes, 803 training prompts across 16 decision points per episode
+  - [x] Full message content, drift flags, superseded markers, dependency info
+  - [x] Prompts include urgency grouping, deadline warnings, stale markers
+- [x] Connect to Unsloth GRPO: notebook loads episodes.json, builds HF Dataset, configures GRPOTrainer
+- [ ] Run training on Colab with GPU and capture reward curves
+- [ ] Log reward components separately to identify which signals are working
+
+### Checkpoint — 6:00 PM (Dinner)
+
+- [x] ✅ Environment generates realistic message streams (76 messages, 19 senders)
+- [x] ✅ Schema drift events fire correctly (tested: superseded messages marked, drift rewards working)
+- [x] ✅ Training notebook configured with GRPO, reward function, and evaluation
+- [ ] ✅ First reward curves from actual GPU training run
+- [ ] 🚨 If messages aren't generating: reduce to 30 messages, fewer senders
+- [ ] 🚨 If training isn't connecting: hardcode environment responses, focus on reward signal
+
+---
+
+## PHASE 3: Integration & Iteration (6:00 PM – 10:00 PM Saturday)
+
+### Person A — Environment Polish
+
+- [x] Polish message content for realism and emotional impact (20+ messages rewritten)
+  - [x] Mom: panicked texting, crying voicemails, medical anxiety, guilt-trip to Tahoe
+  - [x] Boss (Greg): passive-aggressive emails with signature, softens after Meridian postpones
+  - [x] Sister: desperation about kids, voice cracking in voicemail, genuine gratitude
+  - [x] Neighbor Dave: casual bro tone, guilt about Whiskers, neighborhood solidarity
+  - [x] Emma (niece, age 7): 3 kid-perspective messages — pillow fights, scared in the dark, rainbow drawing
+  - [x] FEMA/NWS: kept formal and information-dense (already good)
+- [x] Test all 5 drift events fire correctly and change environment state
+- [x] Verify dependency chains work (23 messages with dependencies, gated in respond_to_message)
+- [x] Edge case handling: stale info penalized (-50% reward), expired deadlines tracked in get_status
+- [x] Redeploy updated environment to HF Spaces (73 messages, timeline engine, drift events)
+- [x] Test HF Spaces deployment works end-to-end remotely (verified via client)
+
+### Person B — Training Optimization
+
+- [ ] Analyze initial reward curves — identify flat or noisy components
+- [ ] Tune reward weights: if one signal dominates, rebalance
+- [ ] If training is flat overall:
+  - [ ] Simplify action space (reduce to 4 actions)
+  - [ ] Increase reward magnitudes
+  - [ ] Reduce episode length
+- [ ] Log specific before/after agent behaviors:
+  - [ ] Capture untrained agent: responds in order, ignores evacuation alert
+  - [ ] Capture trained agent: triages safety first, adapts to drift
+- [ ] Save training checkpoints and reward curve data for demo
+- [ ] Target: 200-500 episodes with visible upward trend
+
+### Checkpoint — 10:00 PM (Doors Close)
+
+- [x] ✅ Environment fully functional with drift events on HF Spaces
+- [ ] ✅ Training curves show upward trend (pending GPU run)
+- [ ] ✅ At least 2 clear before/after behavior examples captured
+- [ ] ✅ All code pushed to GitHub
+- [ ] 🚨 If reward curves are flat: simplify environment overnight, retrain Sunday AM
+- [ ] 🚨 If HF Spaces is broken: prepare local demo as backup
+- [ ] Both teammates agree on Sunday morning priority list before leaving
+
+---
+
+## PHASE 4: Sunday Polish (9:00 AM – 12:00 PM Sunday)
+
+### Person A — Demo & Presentation
+
+- [x] Build demo display (`demo.py`): terminal-based visualization
+  - [x] Color-coded urgency levels (red bg = critical, red = high, yellow = medium, green = low)
+  - [x] Schema drift notifications with magenta banner
+  - [x] Agent action visualization with blue banner and reward display
+  - [x] Two strategies: smart triage vs naive (arrival order)
+  - [x] Comparison mode shows 55% improvement (157.8 vs 101.8 pts)
+  - [x] Coverage breakdown by urgency with ASCII progress bars
+- [x] HF Spaces deployment live with polished environment
+- [x] Write repo README:
+  - [x] Scenario hook, three layers of difficulty, sender profiles table
+  - [x] MCP tools table, reward function table, episode variation details
+  - [x] Quick start (hosted + local), repo structure, tech stack
+- [ ] Draft the 3-minute pitch outline:
+  - [ ] 0:00-0:30 — The scenario hook ("A wildfire just hit. You have 47 unread messages.")
+  - [ ] 0:30-1:30 — Show the environment: message stream, drift events, conflicting tasks
+  - [ ] 1:30-2:15 — Untrained vs trained agent comparison
+  - [ ] 2:15-2:45 — Reward curves and training results
+  - [ ] 2:45-3:00 — Why this matters (real-world impact close)
+
+### Person B — Training & Artifacts
+
+- [ ] Run final training session for cleanest possible reward curves
+- [ ] Export reward curve plots (all 5 components + composite)
+- [ ] Export 3 specific before/after examples with clear narrative:
+  - [ ] Example 1: Untrained ignores FEMA evacuation → Trained prioritizes it first
+  - [ ] Example 2: Untrained misses insurance deadline after policy change → Trained adapts
+  - [ ] Example 3: Untrained sends form-letter reply to Mom → Trained matches emotional tone
+- [x] Finalize Colab notebook:
+  - [x] Clean, commented code with markdown sections
+  - [x] Reward function with test cases (good action vs bad action vs junk)
+  - [x] Evaluation cell comparing trained model choices
+  - [ ] Reward curves visible when run (pending GPU run)
+- [ ] Verify all required artifacts exist:
+  - [ ] Public GitHub repo
+  - [ ] HF Spaces deployment
+  - [ ] Colab training notebook
+  - [ ] YouTube video (not yet recorded)
+
+### Checkpoint — 12:00 PM
+
+- [x] ✅ Demo runs smoothly end-to-end (smart vs naive comparison, 55% improvement)
+- [ ] ✅ Reward curves are clean and trend upward (pending GPU run)
+- [ ] ✅ Before/after examples are compelling (pending GPU run)
+- [x] ✅ README is complete
+- [x] ✅ Colab notebook is clean (14 cells, markdown sections, reward function tests)
+- [ ] 🚨 If demo is buggy: simplify to scripted walkthrough of one episode
+- [ ] 🚨 If reward curves are still flat: show qualitative behavior improvement instead
+
+---
+
+## PHASE 5: Record & Submit (12:00 PM – 1:00 PM Sunday)
+
+### Both Together
+
+- [ ] Record one-minute YouTube video:
+  - [ ] Person A: voiceover narration
+  - [ ] Person B: screen capture / screen share
+  - [ ] Structure:
+    - [ ] 0:00-0:15 — "CrisisInbox: training AI to triage personal tasks during disasters"
+    - [ ] 0:15-0:35 — Show environment: messages arriving, drift event firing
+    - [ ] 0:35-0:50 — Before/after: untrained fails → trained succeeds
+    - [ ] 0:50-1:00 — Flash reward curves, close with impact statement
+  - [ ] Upload to YouTube (unlisted is fine)
+- [ ] Rehearse 3-minute live pitch at least twice out loud
+- [ ] Prepare for Q&A — practice answers to:
+  - [ ] "Why is this hard for LLMs?"
+  - [ ] "How does schema drift work specifically?"
+  - [ ] "What would you build next with more time?"
+  - [ ] "How is this different from a standard chatbot?"
+  - [ ] "How does the reward function handle edge cases?"
+- [ ] **SUBMIT via Cerebral Valley form:**
+  - [ ] GitHub repo URL (public)
+  - [ ] HF Spaces URL
+  - [ ] YouTube video URL
+  - [ ] Colab notebook URL
+  - [ ] Select problem statement 3.2 + Patronus AI sub-theme
+  - [ ] Confirm submission received
+
+---
+
+## FINAL CHECKLIST — Before Judging Starts (1:15 PM Sunday)
+
+- [ ] GitHub repo is public and all code is pushed
+- [ ] HF Spaces deployment is live and accessible
+- [ ] Colab notebook runs end-to-end
+- [ ] YouTube video is uploaded and link works
+- [ ] Submission form completed
+- [ ] Both teammates can explain every part of the project
+- [ ] Laptop charged for demo
+- [ ] Demo runs without internet (backup plan if WiFi fails)
+
+---
+
+## EMERGENCY FALLBACKS
+
+| Problem | Fallback |
+|---------|----------|
+| OpenEnv won't work | Build gym-style env, wrap in OpenEnv interface last |
+| HF Spaces deploy fails | Run demo locally, screenshot HF attempt for judges |
+| Training won't converge | Simplify to 3 actions, 20 messages, 2 drift events |
+| Reward curves flat | Show qualitative behavior change, explain reward design |
+| YouTube upload fails | Screen-record on phone, upload from mobile |
+| One teammate burns out | Other covers — both should understand full stack |
+
+---
+
+## KEY RULE
+
+**If something isn't working after 45 minutes, simplify. A working simple version always beats a broken ambitious one.**

debug_env.py

@@ -0,0 +1,38 @@
+"""Debug script to inspect environment response format."""
+import sys
+sys.path.insert(0, r'C:\Users\smrit\CascadeProjects\crisis-inbox-test')
+
+from crisis_inbox import CrisisInboxEnv
+
+with CrisisInboxEnv(base_url="https://eptan-crisis-inbox.hf.space") as env:
+    env.reset()
+    
+    # List available tools
+    tools = env.list_tools()
+    print(f"Available tools ({len(tools)}):")
+    for t in tools:
+        print(f"  - {t.name}: {t.description[:60]}...")
+    
+    # Get inbox and inspect
+    print("\n--- Calling get_inbox ---")
+    inbox = env.call_tool("get_inbox")
+    print(f"Type: {type(inbox)}")
+    print(f"Content preview: {str(inbox)[:500]}")
+    
+    # Try to get a specific message
+    print("\n--- Trying read_message ---")
+    try:
+        msg = env.call_tool("read_message", message_id="msg_001")
+        print(f"Type: {type(msg)}")
+        print(f"Content: {str(msg)[:300]}")
+    except Exception as e:
+        print(f"Error: {e}")
+    
+    # Try status
+    print("\n--- Calling get_status ---")
+    try:
+        status = env.call_tool("get_status")
+        print(f"Type: {type(status)}")
+        print(f"Content: {str(status)[:300]}")
+    except Exception as e:
+        print(f"Error: {e}")

demo.py

@@ -0,0 +1,401 @@
+#!/usr/bin/env python3
+"""
+CrisisInbox Demo Display
+
+Runs a full episode with color-coded urgency levels, drift event
+notifications, and agent action visualization. Designed for the
+3-minute hackathon demo.
+
+Usage:
+    python demo.py                    # Run against local server
+    python demo.py --remote           # Run against HF Spaces
+    python demo.py --strategy smart   # Use smart triage (default)
+    python demo.py --strategy naive   # Use naive arrival-order strategy
+"""
+
+import argparse
+import json
+import sys
+import time
+
+from server.crisis_inbox_environment import CrisisInboxEnvironment
+from openenv.core.env_server.mcp_types import CallToolAction
+
+
+# ---------------------------------------------------------------------------
+# ANSI colors
+# ---------------------------------------------------------------------------
+class C:
+    RED = "\033[91m"
+    YELLOW = "\033[93m"
+    GREEN = "\033[92m"
+    BLUE = "\033[94m"
+    MAGENTA = "\033[95m"
+    CYAN = "\033[96m"
+    WHITE = "\033[97m"
+    GRAY = "\033[90m"
+    BOLD = "\033[1m"
+    DIM = "\033[2m"
+    UNDERLINE = "\033[4m"
+    RESET = "\033[0m"
+    BG_RED = "\033[41m"
+    BG_YELLOW = "\033[43m"
+    BG_BLUE = "\033[44m"
+    BG_MAGENTA = "\033[45m"
+
+
+URGENCY_COLORS = {
+    "critical": C.BG_RED + C.WHITE + C.BOLD,
+    "high": C.RED + C.BOLD,
+    "medium": C.YELLOW,
+    "low": C.GREEN,
+}
+
+URGENCY_LABELS = {
+    "critical": " CRITICAL ",
+    "high": "   HIGH   ",
+    "medium": "  MEDIUM  ",
+    "low": "   LOW    ",
+}
+
+CHANNEL_ICONS = {
+    "sms": "SMS",
+    "email": "EMAIL",
+    "phone": "PHONE",
+    "government_alert": "ALERT",
+    "app_notification": "APP",
+    "social_media": "SOCIAL",
+}
+
+
+# ---------------------------------------------------------------------------
+# Display helpers
+# ---------------------------------------------------------------------------
+def header(text: str):
+    width = 72
+    print(f"\n{C.BOLD}{C.CYAN}{'=' * width}{C.RESET}")
+    print(f"{C.BOLD}{C.CYAN}  {text}{C.RESET}")
+    print(f"{C.BOLD}{C.CYAN}{'=' * width}{C.RESET}")
+
+
+def subheader(text: str):
+    print(f"\n{C.BOLD}{C.WHITE}--- {text} ---{C.RESET}")
+
+
+def show_time(hour: float):
+    h = int(hour)
+    m = int((hour - h) * 60)
+    day = "Day 1" if hour < 24 else "Day 2"
+    display_h = h % 24
+    ampm = "AM" if display_h < 12 else "PM"
+    display_h = display_h % 12 or 12
+    return f"{day} {display_h}:{m:02d} {ampm} (hour {hour:.1f}/48)"
+
+
+def show_message_line(msg: dict, prefix: str = ""):
+    urg = msg["urgency"]
+    color = URGENCY_COLORS.get(urg, "")
+    label = URGENCY_LABELS.get(urg, urg.upper())
+    channel = CHANNEL_ICONS.get(msg["channel"], msg["channel"])
+    handled = f"{C.GRAY}[DONE]{C.RESET} " if msg.get("handled") else ""
+    superseded = f"{C.GRAY}[STALE]{C.RESET} " if msg.get("superseded") else ""
+    drift = f"{C.BG_MAGENTA}{C.WHITE} DRIFT {C.RESET} " if msg.get("drift_flag") else ""
+    deadline_str = ""
+    if msg.get("deadline_hours") is not None:
+        deadline_str = f" {C.DIM}(due h{msg['deadline_hours']:.0f}){C.RESET}"
+
+    print(
+        f"  {prefix}{color}[{label}]{C.RESET} "
+        f"{C.DIM}{channel:>5}{C.RESET} "
+        f"{drift}{superseded}{handled}"
+        f"{C.BOLD}{msg['sender']}{C.RESET}: {msg['subject']}"
+        f"{deadline_str}"
+    )
+
+
+def show_action(action_text: str):
+    print(f"\n  {C.BG_BLUE}{C.WHITE}{C.BOLD} AGENT ACTION {C.RESET} {C.CYAN}{action_text}{C.RESET}")
+
+
+def show_reward(reward: float, total: float):
+    color = C.GREEN if reward > 0 else C.RED
+    print(f"  {color}+{reward:.1f} pts{C.RESET} {C.DIM}(total: {total:.1f}){C.RESET}")
+
+
+def show_drift_alert(msg: dict):
+    print(f"\n  {C.BG_MAGENTA}{C.WHITE}{C.BOLD} SCHEMA DRIFT {C.RESET} "
+          f"{C.MAGENTA}{C.BOLD}{msg['sender']}: {msg['subject']}{C.RESET}")
+    print(f"  {C.MAGENTA}Rules have changed! Previous information may be outdated.{C.RESET}")
+
+
+def show_expired(expired: list):
+    if expired:
+        print(f"\n  {C.RED}{C.BOLD}EXPIRED DEADLINES:{C.RESET}")
+        for e in expired:
+            print(f"    {C.RED}x {e['subject']}{C.RESET}")
+
+
+def show_upcoming(upcoming: list):
+    if upcoming:
+        print(f"\n  {C.YELLOW}UPCOMING DEADLINES:{C.RESET}")
+        for u in upcoming:
+            print(f"    {C.YELLOW}! {u['subject']} ({u['hours_remaining']:.1f}h left){C.RESET}")
+
+
+def pause(seconds: float = 0.5):
+    time.sleep(seconds)
+
+
+# ---------------------------------------------------------------------------
+# Agent strategies
+# ---------------------------------------------------------------------------
+def smart_priority(messages: list[dict]) -> list[dict]:
+    """Triage: safety first, then deadlines, then drift, then urgency."""
+    urgency_order = {"critical": 0, "high": 1, "medium": 2, "low": 3}
+    unhandled = [m for m in messages if not m.get("handled") and not m.get("superseded")]
+
+    def score(m):
+        urg = urgency_order.get(m["urgency"], 4)
+        deadline = m.get("deadline_hours") or 999
+        drift = 0 if m.get("drift_flag") else 1
+        return (urg, drift, deadline)
+
+    return sorted(unhandled, key=score)
+
+
+def naive_order(messages: list[dict]) -> list[dict]:
+    """Respond in arrival order (bad strategy for comparison)."""
+    return [m for m in messages if not m.get("handled") and not m.get("superseded")]
+
+
+# ---------------------------------------------------------------------------
+# Main demo loop
+# ---------------------------------------------------------------------------
+def run_demo(strategy: str = "smart", seed: int = 42, speed: float = 0.3):
+    env = CrisisInboxEnvironment()
+    obs = env.reset(seed=seed)
+
+    def call(tool, **kwargs):
+        o = env.step(CallToolAction(type="call_tool", tool_name=tool, arguments=kwargs))
+        return json.loads(o.result.data)
+
+    prioritize = smart_priority if strategy == "smart" else naive_order
+    strategy_name = "SMART TRIAGE" if strategy == "smart" else "NAIVE (arrival order)"
+
+    header(f"CrisisInbox Demo  |  Strategy: {strategy_name}")
+    print(f"\n  {C.DIM}Scenario: Post-hurricane evacuation in Sacramento")
+    print(f"  You are a working parent. Your phone is about to explode.{C.RESET}")
+    print(f"  {C.DIM}Messages: {obs.metadata['messages_total']} total | "
+          f"Drift events: {obs.metadata['drift_events_scheduled']} scheduled{C.RESET}")
+    pause(speed * 2)
+
+    last_drift_count = 0
+    actions_taken = 0
+    max_actions = 25  # Cap for demo length
+
+    # Advance through the timeline in chunks
+    time_chunks = [0, 2, 6, 12, 20, 25, 34, 44, 48]
+
+    for i in range(len(time_chunks) - 1):
+        target_hour = time_chunks[i + 1]
+        current = time_chunks[i]
+
+        # Advance time to target
+        while current < target_hour:
+            advance = min(4.0, target_hour - current)
+            result = call("advance_time", hours=advance)
+            current = result["current_hour"]
+            if result.get("new_messages", 0) > 0:
+                pass  # Messages delivered silently
+
+        # Get current state
+        status = call("get_status")
+        inbox = call("get_inbox")
+
+        subheader(f"HOUR {status['current_hour']:.1f}  |  {show_time(status['current_hour'])}")
+        print(f"  {C.DIM}Messages arrived: {status['messages_total_arrived']} | "
+              f"Handled: {status['messages_handled']} | "
+              f"Score: {status['total_score']:.1f}{C.RESET}")
+
+        # Check for new drift events
+        if status["drift_events_fired"] > last_drift_count:
+            drift_msgs = [m for m in inbox if m.get("drift_flag") and not m.get("handled")]
+            for dm in drift_msgs:
+                show_drift_alert(dm)
+                pause(speed)
+            last_drift_count = status["drift_events_fired"]
+
+        # Show deadlines
+        show_expired(status.get("expired_deadlines", []))
+        show_upcoming(status.get("upcoming_deadlines", []))
+
+        # Show current inbox
+        print(f"\n  {C.BOLD}INBOX ({len(inbox)} messages):{C.RESET}")
+        # Show only recent unhandled, plus any drift
+        visible = [m for m in inbox if not m.get("handled")]
+        for msg in visible[:8]:  # Show max 8 at a time
+            show_message_line(msg)
+        if len(visible) > 8:
+            print(f"  {C.DIM}  ... and {len(visible) - 8} more unread{C.RESET}")
+        pause(speed)
+
+        # Agent takes actions on available messages
+        prioritized = prioritize(inbox)
+        actions_this_chunk = 0
+        max_per_chunk = 4  # Don't monopolize one time period
+
+        for msg in prioritized:
+            if actions_taken >= max_actions or actions_this_chunk >= max_per_chunk:
+                break
+
+            msg_id = msg["id"]
+
+            # Read the message first
+            full_msg = call("read_message", message_id=msg_id)
+            if "error" in full_msg:
+                continue
+
+            # Generate a contextual response based on sender/urgency
+            response = _generate_response(full_msg)
+
+            # Try to respond
+            result = call("respond_to_message", message_id=msg_id, response=response)
+            if "error" in result:
+                if "dependencies" in result.get("error", ""):
+                    show_action(f"Cannot handle {msg_id} yet - dependencies unmet")
+                    pause(speed * 0.5)
+                continue
+
+            show_action(f"Respond to {msg['sender']} - \"{msg['subject']}\"")
+            print(f"  {C.DIM}\"{response[:80]}{'...' if len(response) > 80 else ''}\"{C.RESET}")
+            show_reward(result["reward"], result["total_score"])
+            pause(speed)
+
+            actions_taken += 1
+            actions_this_chunk += 1
+
+        if status.get("done"):
+            break
+
+    # Final summary
+    final_status = call("get_status")
+    final_inbox = call("get_inbox")
+
+    header("EPISODE COMPLETE")
+    handled_count = final_status["messages_handled"]
+    total_arrived = final_status["messages_total_arrived"]
+    missed = len(final_status.get("expired_deadlines", []))
+
+    print(f"\n  Strategy: {C.BOLD}{strategy_name}{C.RESET}")
+    print(f"  Final Score: {C.BOLD}{C.GREEN}{final_status['total_score']:.1f} pts{C.RESET}")
+    print(f"  Messages Handled: {handled_count}/{total_arrived}")
+    print(f"  Deadlines Missed: {C.RED}{missed}{C.RESET}")
+    print(f"  Drift Events Encountered: {final_status['drift_events_fired']}")
+
+    # Show what was handled vs missed by urgency
+    handled_ids = {m["id"] for m in final_inbox if m.get("handled")}
+    by_urgency = {"critical": [0, 0], "high": [0, 0], "medium": [0, 0], "low": [0, 0]}
+    for m in final_inbox:
+        urg = m["urgency"]
+        if urg in by_urgency:
+            by_urgency[urg][1] += 1
+            if m["id"] in handled_ids:
+                by_urgency[urg][0] += 1
+
+    print(f"\n  {C.BOLD}Coverage by Urgency:{C.RESET}")
+    for urg, (handled, total) in by_urgency.items():
+        color = URGENCY_COLORS.get(urg, "")
+        pct = (handled / total * 100) if total > 0 else 0
+        bar = "#" * int(pct / 5) + "." * (20 - int(pct / 5))
+        print(f"    {color}{urg:>8}{C.RESET}: {handled}/{total} [{bar}] {pct:.0f}%")
+
+    print()
+    return final_status["total_score"]
+
+
+def _generate_response(msg: dict) -> str:
+    """Generate a contextual response based on sender and content."""
+    sender = msg.get("sender", "")
+    urgency = msg.get("urgency", "")
+    subject = msg.get("subject", "").lower()
+
+    if "national weather" in sender.lower() or "fema" in sender.lower():
+        return "Acknowledged. Following evacuation orders immediately. Heading to designated shelter with essential documents and medication."
+    if sender == "Mom":
+        return "I'm safe, Mom. Don't worry. I'm following the evacuation orders and heading to the shelter. I'll call you as soon as I can. Love you."
+    if sender == "Sister":
+        if "kids" in subject or "take" in subject:
+            return "I'll get them, don't worry. Heading to Oakwood now. I'll text you the second I have them safe."
+        return "Kids are safe with me. Emma is being brave. Don't worry about anything, just stay safe yourself."
+    if "emma" in sender.lower():
+        return "Hey sweetie! Mac and cheese sounds perfect. Mommy is coming soon. Let's build a blanket fort while we wait!"
+    if sender == "Boss" or sender == "HR Department":
+        if "drift" in str(msg.get("drift_flag", "")):
+            return "Thanks for the update on the emergency leave policy. I've submitted my status form on the HR portal. Will be taking the emergency leave days."
+        return "I'm in the evacuation zone and following mandatory orders. Will work remotely when I can access wifi. Updating my status on the portal now."
+    if "insurance" in sender.lower() or "state farm" in sender.lower():
+        return "Filing claim now with policy number and damage photos. Have documented all damage with timestamps before any cleanup."
+    if sender == "Neighbor Dave":
+        return "Thanks for the heads up Dave. Stay safe at the shelter. I'll keep an eye on things here. We'll get through this."
+    if "delta" in sender.lower() or "airline" in sender.lower():
+        return "Selecting Option A for rebooking to the earliest available flight. Thank you for the flexibility during this emergency."
+    if "school" in sender.lower() or "oakwood" in sender.lower():
+        return "Acknowledged. Will arrange pickup before the deadline. Thank you for the early notification."
+    if "pharmacy" in sender.lower() or "cvs" in sender.lower():
+        return "Will pick up the prescription today. If the usual location is closed, I'll transfer to the nearest open CVS."
+    if "sacramento" in sender.lower():
+        return "Acknowledged. Following all advisories and avoiding affected areas."
+    if urgency == "critical":
+        return "Taking immediate action on this critical matter. Will follow up with details shortly."
+    if urgency == "high":
+        return "Understood, this is a priority. Handling this now and will confirm when complete."
+    return "Thank you for the information. I've noted this and will address it as soon as possible given the current situation."
+
+
+# ---------------------------------------------------------------------------
+# Comparison mode
+# ---------------------------------------------------------------------------
+def run_comparison(seed: int = 42, speed: float = 0.2):
+    """Run both strategies side by side for the demo."""
+    header("CRISISINBOX: STRATEGY COMPARISON")
+    print(f"\n  {C.DIM}Same episode, two different approaches.{C.RESET}")
+    print(f"  {C.DIM}Which agent handles a disaster better?{C.RESET}\n")
+    pause(1)
+
+    print(f"{C.RED}{C.BOLD}{'=' * 72}")
+    print(f"  ROUND 1: NAIVE AGENT (responds in arrival order)")
+    print(f"{'=' * 72}{C.RESET}")
+    pause(0.5)
+    naive_score = run_demo(strategy="naive", seed=seed, speed=speed)
+
+    pause(1)
+    print(f"\n{C.GREEN}{C.BOLD}{'=' * 72}")
+    print(f"  ROUND 2: TRAINED AGENT (smart triage)")
+    print(f"{'=' * 72}{C.RESET}")
+    pause(0.5)
+    smart_score = run_demo(strategy="smart", seed=seed, speed=speed)
+
+    header("FINAL COMPARISON")
+    improvement = smart_score - naive_score
+    pct = (improvement / max(naive_score, 1)) * 100
+    print(f"\n  Naive Agent:   {C.RED}{naive_score:.1f} pts{C.RESET}")
+    print(f"  Trained Agent: {C.GREEN}{smart_score:.1f} pts{C.RESET}")
+    print(f"  Improvement:   {C.BOLD}{C.CYAN}+{improvement:.1f} pts ({pct:.0f}%){C.RESET}\n")
+
+
+# ---------------------------------------------------------------------------
+# Entry point
+# ---------------------------------------------------------------------------
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="CrisisInbox Demo")
+    parser.add_argument("--strategy", choices=["smart", "naive", "compare"],
+                        default="compare", help="Agent strategy (default: compare)")
+    parser.add_argument("--seed", type=int, default=42, help="Episode seed")
+    parser.add_argument("--speed", type=float, default=0.3,
+                        help="Pause between actions in seconds (default: 0.3)")
+    args = parser.parse_args()
+
+    if args.strategy == "compare":
+        run_comparison(seed=args.seed, speed=args.speed)
+    else:
+        run_demo(strategy=args.strategy, seed=args.seed, speed=args.speed)

drift_events.py

@@ -86,9 +86,10 @@ DRIFT_EVACUATION_EXPANSION = DriftEvent(
             channel=Channel.SMS,
             subject="Office is in Zone B - disregard earlier email",
             content=(
-                "FYI the office is in Zone B which just got the evacuation order. Obviously "
-                "don't come in. Work from wherever. But I STILL need those Meridian slides. "
-                "Can you do them from the shelter?"
+                "Well, our office is in Zone B. So obviously disregard the 'come in tomorrow' "
+                "thing. But I STILL need those slides. Can you work on them from wherever "
+                "you are? Sarah said she can pick up some of the slack but I need to know "
+                "what you can deliver. This client won't wait for a hurricane."
             ),
             urgency=Urgency.MEDIUM,
             timestamp_hours=21.5,

generate_episodes.py

@@ -0,0 +1,305 @@
+"""
+Episode Generator for CrisisInbox GRPO Training.
+
+Generates training episodes locally (no server needed) by simulating the
+environment and capturing inbox snapshots at key decision points.
+
+Each episode produces multiple training prompts — one per decision point —
+where the model must choose which message to handle next.
+
+Output format (per episode):
+{
+  "episode_id": "ep_000",
+  "seed": 42,
+  "total_messages": 73,
+  "drift_events": ["drift_insurance", "drift_evacuation", "drift_fema"],
+  "decision_points": [
+    {
+      "hour": 0.0,
+      "prompt": "...",           # full text prompt for the LLM
+      "visible_messages": [...], # inbox snapshot
+      "handled_ids": [],
+      "pending_deadlines": [...],
+      "drift_events_fired": []
+    },
+    ...
+  ]
+}
+"""
+
+import json
+import random
+from typing import Any
+
+from models import Channel, Message, Urgency
+from messages import ALL_MESSAGES
+from drift_events import ALL_DRIFT_EVENTS, DriftEvent, select_drift_events
+
+
+SYSTEM_PROMPT = """You are managing a personal crisis inbox during a post-hurricane evacuation in Sacramento. You are a working parent with 48 hours to triage incoming messages from family, employer, government, insurance, and service providers.
+
+Rules:
+- Reading a message costs 0.1 hours (6 minutes)
+- Responding to a message costs 0.25 hours (15 minutes)
+- You cannot handle everything — prioritize wisely
+- Safety-critical messages (evacuations, medical) should come first
+- Watch for policy changes that supersede earlier information
+- Some messages have dependencies that must be handled first
+- Deadlines are real — missing them reduces your score
+
+Available actions:
+- respond_to_message(message_id, response) — handle a message
+- advance_time(hours) — skip forward to see new messages
+- get_status() — check time, score, deadlines"""
+
+
+def build_episode(seed: int) -> dict[str, Any]:
+    """Build a single training episode with decision-point snapshots."""
+    rng = random.Random(seed)
+
+    # Select drift events
+    drift_events = select_drift_events(count=3, rng=rng)
+    drift_event_ids = [d.id for d in drift_events]
+
+    # Collect drift message IDs
+    selected_drift_msg_ids = set()
+    for drift in drift_events:
+        for msg in drift.messages:
+            selected_drift_msg_ids.add(msg.id)
+
+    all_drift_msg_ids = set()
+    for drift in ALL_DRIFT_EVENTS:
+        for msg in drift.messages:
+            all_drift_msg_ids.add(msg.id)
+
+    # Build message pool with jitter
+    all_messages = []
+    for msg in ALL_MESSAGES:
+        if msg.id in all_drift_msg_ids and msg.id not in selected_drift_msg_ids:
+            continue
+        m = msg.model_copy()
+        if m.timestamp_hours > 0:
+            jitter = rng.uniform(-0.15, 0.15) * m.timestamp_hours
+            m.timestamp_hours = round(max(0.1, min(47.5, m.timestamp_hours + jitter)), 2)
+        if m.deadline_hours is not None and m.deadline_hours > 0:
+            d_jitter = rng.uniform(-0.1, 0.1) * m.deadline_hours
+            m.deadline_hours = round(max(m.timestamp_hours + 0.5, min(72.0, m.deadline_hours + d_jitter)), 2)
+        all_messages.append(m)
+
+    # Also add drift event messages
+    for drift in drift_events:
+        for msg in drift.messages:
+            if not any(m.id == msg.id for m in all_messages):
+                all_messages.append(msg)
+
+    # Sort all messages by arrival time
+    all_messages.sort(key=lambda m: m.timestamp_hours)
+
+    # Track superseded messages
+    superseded = {}
+    for drift in drift_events:
+        for old_id in drift.superseded_msg_ids:
+            for dmsg in drift.messages:
+                if dmsg.supersedes == old_id:
+                    superseded[old_id] = dmsg.id
+
+    # Simulate the episode at key time points to capture decision snapshots
+    decision_hours = [0.0, 2.0, 6.0, 10.0, 14.0, 18.0]
+    # Add drift trigger hours
+    for drift in drift_events:
+        decision_hours.append(drift.trigger_hour)
+        decision_hours.append(drift.trigger_hour + 1.0)
+    # Add late-game hours
+    decision_hours.extend([28.0, 34.0, 40.0, 44.0, 47.0])
+    decision_hours = sorted(set(decision_hours))
+
+    decision_points = []
+    fired_drifts = set()
+
+    for hour in decision_hours:
+        if hour > 48.0:
+            continue
+
+        # Deliver messages visible at this hour
+        visible = [m for m in all_messages if m.timestamp_hours <= hour]
+
+        # Fire drift events
+        newly_fired = []
+        for drift in drift_events:
+            if drift.id not in fired_drifts and hour >= drift.trigger_hour:
+                fired_drifts.add(drift.id)
+                newly_fired.append(drift.id)
+
+        # Build inbox summary
+        visible_summaries = []
+        for msg in visible:
+            is_superseded = msg.id in superseded
+            summary = {
+                "id": msg.id,
+                "sender": msg.sender,
+                "subject": msg.subject,
+                "content": msg.content,
+                "urgency": msg.urgency.value,
+                "channel": msg.channel.value,
+                "timestamp_hours": msg.timestamp_hours,
+                "deadline_hours": msg.deadline_hours,
+                "dependencies": msg.dependencies,
+                "drift_flag": msg.drift_flag,
+                "superseded": is_superseded,
+            }
+            visible_summaries.append(summary)
+
+        # Identify pending deadlines
+        pending_deadlines = []
+        for msg in visible:
+            if msg.deadline_hours is not None:
+                remaining = msg.deadline_hours - hour
+                if remaining > 0:
+                    pending_deadlines.append({
+                        "id": msg.id,
+                        "subject": msg.subject,
+                        "urgency": msg.urgency.value,
+                        "hours_remaining": round(remaining, 1),
+                    })
+                elif remaining > -2:  # recently expired
+                    pending_deadlines.append({
+                        "id": msg.id,
+                        "subject": msg.subject,
+                        "urgency": msg.urgency.value,
+                        "hours_remaining": round(remaining, 1),
+                        "expired": True,
+                    })
+
+        # Build the text prompt
+        prompt = format_prompt(hour, visible_summaries, pending_deadlines, list(fired_drifts))
+
+        decision_points.append({
+            "hour": hour,
+            "prompt": prompt,
+            "visible_count": len(visible),
+            "visible_messages": visible_summaries,
+            "pending_deadlines": pending_deadlines,
+            "drift_events_fired": list(fired_drifts),
+            "newly_fired_drifts": newly_fired,
+        })
+
+    return {
+        "episode_id": f"ep_{seed:03d}",
+        "seed": seed,
+        "total_messages": len(all_messages),
+        "drift_events": drift_event_ids,
+        "superseded_messages": superseded,
+        "decision_points": decision_points,
+    }
+
+
+def format_prompt(hour: float, messages: list, deadlines: list, fired_drifts: list,
+                  max_messages: int = 20) -> str:
+    """Format an inbox state into a text prompt for the LLM.
+
+    Only the top `max_messages` unhandled messages are shown (by urgency then
+    deadline), keeping prompts within ~1500 tokens for small-model training.
+    """
+    lines = [SYSTEM_PROMPT, ""]
+    lines.append(f"CURRENT TIME: Hour {hour:.1f} of 48 ({48 - hour:.1f} hours remaining)")
+    lines.append(f"MESSAGES IN INBOX: {len(messages)}")
+    lines.append("")
+
+    # Show urgent deadlines first
+    urgent = [d for d in deadlines if not d.get("expired") and d["hours_remaining"] < 4]
+    expired = [d for d in deadlines if d.get("expired")]
+    if urgent:
+        lines.append("URGENT DEADLINES:")
+        for d in sorted(urgent, key=lambda x: x["hours_remaining"]):
+            lines.append(f"  ! {d['subject']} — {d['hours_remaining']}h left [{d['urgency']}]")
+        lines.append("")
+    if expired:
+        lines.append("EXPIRED DEADLINES:")
+        for d in expired:
+            lines.append(f"  x {d['subject']} — expired {abs(d['hours_remaining']):.1f}h ago")
+        lines.append("")
+
+    if fired_drifts:
+        lines.append(f"POLICY CHANGES DETECTED: {len(fired_drifts)}")
+        lines.append("")
+
+    # Prioritize unhandled messages by urgency then deadline
+    urgency_order = {"critical": 0, "high": 1, "medium": 2, "low": 3}
+    ranked = sorted(messages, key=lambda m: (
+        urgency_order.get(m["urgency"], 4),
+        0 if m.get("drift_flag") else 1,
+        m.get("deadline_hours") or 999,
+    ))
+    shown = ranked[:max_messages]
+    omitted = len(messages) - len(shown)
+
+    # Group shown messages by urgency
+    by_urgency = {"critical": [], "high": [], "medium": [], "low": []}
+    for msg in shown:
+        by_urgency.get(msg["urgency"], by_urgency["low"]).append(msg)
+
+    for level in ["critical", "high", "medium", "low"]:
+        msgs = by_urgency[level]
+        if msgs:
+            lines.append(f"--- {level.upper()} ({len(msgs)}) ---")
+            for msg in msgs:
+                stale = " [STALE]" if msg.get("superseded") else ""
+                drift = " [POLICY CHANGE]" if msg.get("drift_flag") else ""
+                deadline = f" (due h{msg['deadline_hours']})" if msg.get("deadline_hours") else ""
+                deps = f" [requires: {', '.join(msg['dependencies'])}]" if msg.get("dependencies") else ""
+                lines.append(f"  [{msg['id']}] {msg['sender']}: {msg['subject']}{deadline}{stale}{drift}{deps}")
+                # Show content preview (first 120 chars)
+                preview = msg["content"][:120].replace("\n", " ")
+                if len(msg["content"]) > 120:
+                    preview += "..."
+                lines.append(f"    > {preview}")
+            lines.append("")
+
+    if omitted > 0:
+        lines.append(f"  ({omitted} lower-priority messages not shown)")
+        lines.append("")
+
+    lines.append("Which message should you handle next? Respond with respond_to_message(message_id, response).")
+    return "\n".join(lines)
+
+
+def generate_episodes(num_episodes: int = 50, start_seed: int = 1000) -> list:
+    """Generate multiple training episodes with different seeds."""
+    episodes = []
+    for i in range(num_episodes):
+        seed = start_seed + i
+        print(f"  Episode {i + 1}/{num_episodes} (seed={seed})...", end=" ")
+        episode = build_episode(seed)
+        n_dp = len(episode["decision_points"])
+        n_msg = episode["total_messages"]
+        drifts = ", ".join(episode["drift_events"])
+        print(f"{n_msg} messages, {n_dp} decision points, drifts: [{drifts}]")
+        episodes.append(episode)
+    return episodes
+
+
+def save_episodes(episodes: list, filename: str = "episodes.json"):
+    """Save episodes to JSON file."""
+    with open(filename, "w") as f:
+        json.dump(episodes, f, indent=2)
+    total_prompts = sum(len(ep["decision_points"]) for ep in episodes)
+    print(f"\nSaved {len(episodes)} episodes ({total_prompts} training prompts) to {filename}")
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Generate CrisisInbox training episodes")
+    parser.add_argument("-n", "--num-episodes", type=int, default=50, help="Number of episodes")
+    parser.add_argument("-s", "--start-seed", type=int, default=1000, help="Starting seed")
+    parser.add_argument("-o", "--output", type=str, default="episodes.json", help="Output file")
+    parser.add_argument("--sample", type=int, default=5, help="Also save N sample episodes")
+    args = parser.parse_args()
+
+    print(f"Generating {args.num_episodes} episodes (seeds {args.start_seed}-{args.start_seed + args.num_episodes - 1})...")
+    episodes = generate_episodes(args.num_episodes, args.start_seed)
+    save_episodes(episodes, args.output)
+
+    if args.sample > 0:
+        sample_file = "sample_episodes.json"
+        save_episodes(episodes[:args.sample], sample_file)

messages.py

@@ -46,9 +46,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Are you safe?",
         content=(
-            "Honey are you ok?? I saw the news about the hurricane. Your father and I "
-            "are worried sick. Please call us when you can. We can drive down to help "
-            "if you need us."
+            "Honey are you ok?? Just saw channel 3 and they're showing your neighborhood. "
+            "I'm shaking. Your father keeps pacing. PLEASE call us the second you see this. "
+            "We will drive down right now if you need us. I don't care about the roads. "
+            "I love you so much please be safe."
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=0.5,
@@ -59,9 +60,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Can you take the kids?",
         content=(
-            "Hey, my office is making us come in even with the storm. Can you pick up "
-            "Emma and Jake from Oakwood Elementary by 3pm? They're doing early dismissal. "
-            "I know it's a lot but I have no one else to ask."
+            "Hey I need a huge favor. My idiot boss is making us come in even with "
+            "the storm. Can you PLEASE pick up Emma and Jake from Oakwood by 3pm? "
+            "They're doing early dismissal and I literally have no one else. "
+            "Emma gets scared during storms so just tell her it's an adventure ok? "
+            "I owe you big time."
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=0.75,
@@ -88,9 +91,14 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.EMAIL,
         subject="Tomorrow's status - need you in office",
         content=(
-            "Team, I know the storm is bad but we have the Meridian client presentation "
-            "Thursday. I need everyone on deck tomorrow morning. If you absolutely cannot "
-            "make it in, use PTO. No exceptions. Let me know your status by tonight."
+            "Team,\n\n"
+            "I realize the weather situation isn't ideal but the Meridian presentation "
+            "is Thursday and the client isn't going to reschedule. I need everyone in the "
+            "office tomorrow morning, full stop. If you ABSOLUTELY cannot make it, submit "
+            "PTO through the portal. I shouldn't have to remind you this is a career-defining "
+            "account.\n\n"
+            "Confirm your status by tonight.\n\n"
+            "— Greg"
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=1.5,
@@ -117,9 +125,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Need help boarding windows",
         content=(
-            "Hey, my back is out and I can't board up my windows alone. Any chance you "
-            "could help me out for 20 minutes? I've got the plywood and screws ready. "
-            "Happy to return the favor."
+            "Hey it's Dave from next door. My back went out yesterday and I can't "
+            "board up these windows by myself. I know you've got your own stuff going on "
+            "but if you've got 20 min I've got the plywood and drill ready to go. "
+            "I'll owe you a case of beer when this is all over."
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=2.0,
@@ -177,9 +186,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.PHONE,
         subject="Missed call from Mom",
         content=(
-            "You have a missed call from Mom. She left a voicemail: 'Sweetie please call "
-            "me back. Your dad says we should all go to Uncle Rick's in Tahoe. He has room. "
-            "But we need to leave soon. Love you.'"
+            "You have a missed call from Mom. Voicemail: '[sounds like she's been crying] "
+            "Baby please call me back. I can't stop watching the news. Your dad talked to "
+            "Uncle Rick and he says come to Tahoe, he's got plenty of room. But honey we "
+            "need to leave SOON before the roads get worse. Please please call me. I love you.'"
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=4.0,
@@ -190,9 +200,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Water rising on our street",
         content=(
-            "Hey heads up - water is starting to come up on Elm St. My basement is already "
-            "getting wet. You should move anything valuable upstairs ASAP. I'm thinking "
-            "about going to the shelter."
+            "Dude the water is coming up FAST on Elm. My basement already has 6 inches. "
+            "Get anything you care about upstairs RIGHT NOW. I'm not trying to scare you "
+            "but this is way worse than they said it would be. I think I'm heading to "
+            "the shelter. You should too."
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=4.5,
@@ -233,8 +244,9 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Did you get my email?",
         content=(
-            "Haven't heard from you about tomorrow. I need a yes or no on the office. "
-            "The Meridian deck won't finish itself."
+            "Still waiting on your status. Yes or no on tomorrow? I've got half the team "
+            "going dark on me. The Meridian deck won't finish itself and I'm not going to "
+            "be the one explaining to Jacobs why we weren't prepared."
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=6.5,
@@ -247,8 +259,9 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Did you get the kids?",
         content=(
-            "Hey did you manage to get Emma and Jake? The school keeps calling me but "
-            "I'm stuck in a meeting. Please let me know they're safe."
+            "Hey the school keeps calling me and I can't pick up because I'm in this "
+            "stupid meeting. Did you get them?? Emma's teacher said she was crying. "
+            "Please just text me back one word. PLEASE. I'm losing it over here."
         ),
         urgency=Urgency.CRITICAL,
         timestamp_hours=7.0,
@@ -273,9 +286,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Left my cat - can you check?",
         content=(
-            "I evacuated to Lincoln High but I had to leave my cat Whiskers. She's in "
-            "the upstairs bathroom with food and water. If you're still on the street "
-            "can you check on her? I feel terrible."
+            "Man I feel awful. I evacuated to Lincoln High but I had to leave Whiskers "
+            "behind. She's in the upstairs bathroom — I left food and water but she was "
+            "meowing so loud when I shut the door. If you're still on the street can you "
+            "just peek in and make sure she's ok? The spare key is under the gnome."
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=8.0,
@@ -300,9 +314,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Uncle Rick says come to Tahoe",
         content=(
-            "Uncle Rick confirmed you can all stay at his place in Tahoe. He has 3 spare "
-            "bedrooms. Your sister and the kids can come too. Dad and I are already packing. "
-            "Can you be ready by morning? The roads north should still be clear."
+            "OK so Uncle Rick says absolutely yes, come up. He's got 3 bedrooms and says "
+            "stay as long as you need. Your sister and the kids too. Dad and I are packing "
+            "the car now. Honey PLEASE come. I won't be able to sleep tonight if you're "
+            "still down there. The roads north are clear, we checked. Can you be ready by morning?"
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=9.0,
@@ -373,10 +388,12 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.EMAIL,
         subject="RE: Tomorrow's status - updated guidance",
         content=(
-            "Alright team, I spoke with leadership. If you're in an evacuation zone, "
-            "work from wherever you can. But I still need the Meridian deck finalized "
-            "by Thursday 5 PM. Whoever is available, claim your slides in the shared doc. "
-            "No excuses on the deadline."
+            "Team,\n\n"
+            "Spoke with leadership. Fine — if you're in an evac zone, work remote. "
+            "But the Meridian deck is still due Thursday 5 PM, no exceptions. I don't care "
+            "if you're working from a shelter, a Starbucks, or your car. Claim your slides "
+            "in the shared doc by end of day or I'll reassign them.\n\n"
+            "— Greg"
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=12.0,
@@ -398,15 +415,29 @@ ALL_MESSAGES: list[Message] = [
         timestamp_hours=13.0,
         deadline_hours=48.0,
     ),
+    Message(
+        id="msg_026b",
+        sender="Emma (niece)",
+        channel=Channel.SMS,
+        subject="from emmas phone",
+        content=(
+            "hi its emma. jake took the big pillow and wont share. also when is mommy "
+            "coming. are we having a sleepover?? can we have mac and cheese for dinner. "
+            "i miss mr buttons. love you"
+        ),
+        urgency=Urgency.LOW,
+        timestamp_hours=12.5,
+    ),
     Message(
         id="msg_027",
         sender="Sister",
         channel=Channel.SMS,
         subject="Kids are asking about their stuff",
         content=(
-            "Emma keeps crying about her stuffed bunny she left at school. Jake wants his "
-            "tablet for games. Did you grab their backpacks when you picked them up? "
-            "Also my office finally let us go home. I can come get them around 8 PM."
+            "Emma won't stop crying about Mr. Buttons — that's her stuffed bunny she left "
+            "in her cubby at school. Jake is being Jake, just wants his tablet. Did you "
+            "happen to grab their backpacks? My office FINALLY let us leave. I can come "
+            "get them around 8 if the roads are ok. Thank you for doing this. Seriously."
         ),
         urgency=Urgency.LOW,
         timestamp_hours=13.5,
@@ -445,10 +476,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Dad's medication",
         content=(
-            "Honey, quick question - your dad forgot his heart medication at home. "
-            "We're almost to Tahoe. Can you call his doctor tomorrow and see if they "
-            "can call in a prescription to a pharmacy up here? Dr. Patel, number is "
-            "in dad's phone contacts."
+            "Honey I don't want to add to your plate but your dad just realized he left "
+            "his heart medication on the kitchen counter. We're already past Folsom. "
+            "Can you PLEASE call Dr. Patel tomorrow morning and ask them to call in a "
+            "refill to a pharmacy in Truckee? I'm trying not to panic but you know "
+            "he can't miss more than a day. His number should be in dad's phone under P."
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=16.0,
@@ -511,15 +543,30 @@ ALL_MESSAGES: list[Message] = [
         urgency=Urgency.HIGH,
         timestamp_hours=19.0,
     ),
+    Message(
+        id="msg_034b",
+        sender="Emma (niece)",
+        channel=Channel.SMS,
+        subject="im scared",
+        content=(
+            "the lights went off and its really dark. jake is pretending hes not scared "
+            "but he is. when is mommy coming?? you said she was coming. "
+            "i dont like the wind noise. can you come sit with us"
+        ),
+        urgency=Urgency.MEDIUM,
+        timestamp_hours=19.0,
+    ),
     Message(
         id="msg_035",
         sender="Sister",
         channel=Channel.PHONE,
         subject="Missed call from Sister",
         content=(
-            "Missed call from Sister. Voicemail: 'Hey, I'm on my way to get the kids "
-            "but the roads are insane. GPS says Florin Rd is flooded. Can you keep them "
-            "tonight? I'll figure it out in the morning. Sorry to ask.'"
+            "Missed call from Sister. Voicemail: '[car horn honking in background] Hey it's me, "
+            "I'm trying to get to you but Florin is completely underwater and my GPS keeps "
+            "rerouting me in circles. I think I need to turn back. Can you keep them tonight? "
+            "Emma's probably already asleep, just — tell her Mommy loves her. I'm so sorry. "
+            "I'll figure it out in the morning. [voice cracks] Thank you.'"
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=19.5,
@@ -568,9 +615,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Office is in Zone B - disregard earlier email",
         content=(
-            "FYI the office is in Zone B which just got the evacuation order. Obviously "
-            "don't come in. Work from wherever. But I STILL need those Meridian slides. "
-            "Can you do them from the shelter?"
+            "Well, our office is in Zone B. So obviously disregard the 'come in tomorrow' "
+            "thing. But I STILL need those slides. Can you work on them from wherever "
+            "you are? Sarah said she can pick up some of the slack but I need to know "
+            "what you can deliver. This client won't wait for a hurricane."
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=21.5,
@@ -611,10 +659,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Someone broke into houses on our street",
         content=(
-            "Just heard from the cops at the shelter - there have been break-ins on Elm St "
-            "and Oak Ave. Looters taking advantage of evacuations. If you left anything "
-            "valuable, you might want to report it. Cops said to call the non-emergency "
-            "line to file a report: 555-0199."
+            "Hey bad news. Talked to a cop here at the shelter and he said there have been "
+            "break-ins on Elm and Oak. People are the worst — looting during a disaster. "
+            "If you left anything valuable at your place you should probably file a report. "
+            "Non-emergency line is 555-0199. Guy two cots over from me said they got his TV "
+            "and laptop. Unbelievable."
         ),
         urgency=Urgency.HIGH,
         timestamp_hours=23.0,
@@ -702,9 +751,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Did you call Dr. Patel?",
         content=(
-            "Honey did you get a chance to call Dr. Patel about dad's medication? "
-            "He's been without it for a day now. I'm getting worried. The pharmacy "
-            "up here says they need the doctor to call it in."
+            "Sweetheart did you call Dr. Patel?? Dad hasn't had his medication in over "
+            "24 hours now and I can tell his blood pressure is up. He won't admit it "
+            "but he's been dizzy. The pharmacy in Truckee says they need the doctor to "
+            "call it in — they can't just refill it. I know you have a million things "
+            "going on but this one really scares me."
         ),
         urgency=Urgency.CRITICAL,
         timestamp_hours=27.0,
@@ -804,9 +855,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.EMAIL,
         subject="Meridian presentation postponed",
         content=(
-            "Good news (sort of) - Meridian postponed the presentation to next week. "
-            "Half their team is dealing with the storm too. Take the emergency leave, "
-            "focus on your family. We'll regroup Monday."
+            "OK update — Meridian called and pushed the presentation to next Wednesday. "
+            "Turns out half their team is dealing with the storm too. Go figure.\n\n"
+            "Take the emergency leave. Focus on your family. I know I was being intense "
+            "earlier, I'm dealing with this too. We'll regroup Monday.\n\n"
+            "— Greg"
         ),
         urgency=Urgency.LOW,
         timestamp_hours=33.0,
@@ -836,9 +889,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Dad's doing better",
         content=(
-            "Good news - Dr. Patel called in the prescription to a Rite Aid in Truckee. "
-            "Dad picked it up this afternoon. He's feeling better already. When are you "
-            "coming up? We miss you."
+            "Oh thank God — Dr. Patel called in the prescription to the Rite Aid in Truckee "
+            "and dad picked it up an hour ago. He's already looking better, the color is "
+            "back in his face. Thank you honey. Now PLEASE come up here. Rick made up the "
+            "guest room and everything. We miss you. Mom loves you."
         ),
         urgency=Urgency.LOW,
         timestamp_hours=34.5,
@@ -921,9 +975,9 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Going home to check damage",
         content=(
-            "Zone A evac is lifted! I'm going back to check on the house. Want to come? "
-            "I can swing by and pick you up. Cops said the looters have been caught so "
-            "it should be safe. Let me know."
+            "Zone A evac is lifted!! I'm heading back to see how bad it is. Want to come "
+            "with me? Safety in numbers and all that. Cops said they caught the guys who "
+            "were breaking in so it should be safe. I can swing by in 20 if you're in."
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=38.5,
@@ -935,9 +989,10 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="We're at Tahoe",
         content=(
-            "Made it to Uncle Rick's with the kids. They're already playing in the snow "
-            "like nothing happened. Mom made her soup. There's room for you. "
-            "Seriously, come up when you can."
+            "We made it to Uncle Rick's. The kids are already outside throwing snowballs "
+            "like the last 24 hours never happened. Kids are incredible. Mom immediately "
+            "started making chicken soup because of course she did. There's room for you. "
+            "For real, come up when you can. You deserve a break after everything."
         ),
         urgency=Urgency.LOW,
         timestamp_hours=39.0,
@@ -995,9 +1050,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Please come to Tahoe",
         content=(
-            "Sweetheart, we're all here and we're worried about you still being down "
-            "there. The roads are clear now. Dad's feeling much better. Rick is grilling "
-            "tonight. Please come, even if just for a few days."
+            "Honey everyone is here except you and it doesn't feel right. Dad's feeling "
+            "so much better — he's out on the deck with Rick right now. The roads are "
+            "totally clear, I checked three times. Even if it's just for a couple days. "
+            "You've been taking care of everyone else this whole time. Let us take care "
+            "of you for a bit. Please."
         ),
         urgency=Urgency.MEDIUM,
         timestamp_hours=44.0,
@@ -1064,9 +1121,24 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="House update",
         content=(
-            "My house isn't too bad - some water in the basement but the first floor "
-            "is okay. How's yours? Let me know if you need help with cleanup. I owe you "
-            "one for checking on Whiskers."
+            "Good news — my place isn't as bad as I thought. Water in the basement but "
+            "first floor survived. Could've been way worse. How'd yours make out? "
+            "Seriously, anything you need with cleanup just say the word. You checked on "
+            "Whiskers for me and I won't forget that. Neighbors gotta stick together."
+        ),
+        urgency=Urgency.LOW,
+        timestamp_hours=47.0,
+    ),
+    Message(
+        id="msg_072b",
+        sender="Emma (niece)",
+        channel=Channel.SMS,
+        subject="i drew you a pictur",
+        content=(
+            "hi its emma again!! i drew you a rainbow and a house and us. mommy said "
+            "i can send you a foto. jake says thank you for the adventure. i say thank "
+            "you for the mac and cheese. your the best. when can we do a sleepover again "
+            "but WITHOUT the storm part. love emma age 7"
         ),
         urgency=Urgency.LOW,
         timestamp_hours=47.0,
@@ -1077,9 +1149,11 @@ ALL_MESSAGES: list[Message] = [
         channel=Channel.SMS,
         subject="Kids say thank you",
         content=(
-            "Emma drew you a picture of a rainbow. Jake says 'thanks for the adventure.' "
-            "Kids are so resilient. Thank you for everything during the craziness. "
-            "You're the best sibling ever. Love you."
+            "Emma drew you a picture of a rainbow with 'thank you' written in crayon. "
+            "She wants me to send you a photo. And Jake keeps telling everyone about "
+            "'the adventure' and how you let him eat cereal for dinner. Kids man. "
+            "I don't know what I would have done without you this week. I really don't. "
+            "You're my person. Love you."
         ),
         urgency=Urgency.LOW,
         timestamp_hours=47.5,

notebooks/crisisinbox_grpo_simple.ipynb

@@ -0,0 +1,108 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "# CrisisInbox GRPO Training\n\nTrain a small LLM to triage crisis inbox messages using Group Relative Policy Optimization.\n\n**What this does:**\n1. Loads pre-generated episode data (inbox snapshots at decision points)\n2. For each prompt, the model generates an action (which message to handle + response)\n3. A reward function scores the action based on urgency, deadline, drift adaptation\n4. GRPO updates the model to prefer higher-reward actions\n\n**GPU profiles:**\n- **T4 / free Colab**: Qwen2.5-0.5B, 2048 ctx, 4-bit — runs in ~30 min\n- **H100 / A100**: Qwen2.5-3B, 4096 ctx, 4-bit — better quality, ~20 min"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "source": "# Install dependencies\n!pip install unsloth trl transformers datasets accelerate peft -q\n!pip install huggingface_hub -q\n\n# Download episode data\n# Option 1: From HF dataset (recommended)\n# Option 2: From GitHub repo\n# Option 3: Generate locally with `python generate_episodes.py -n 100`\n\nimport os\nif not os.path.exists(\"episodes.json\"):\n    print(\"Downloading episodes.json from GitHub...\")\n    !wget -q --show-progress https://raw.githubusercontent.com/eptan/crisis-inbox/main/episodes.json\n    if not os.path.exists(\"episodes.json\"):\n        print(\"ERROR: Download failed. Upload episodes.json manually or generate with:\")\n        print(\"  !git clone https://github.com/eptan/crisis-inbox.git && cd crisis-inbox && python generate_episodes.py -n 100\")\nelse:\n    print(\"episodes.json already exists, skipping download\")\n\nprint(\"Setup complete\")",
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": "# === GPU PROFILE ===\n# Change this one variable to switch between T4 and H100 configs.\n# Everything else adapts automatically.\n\nimport torch\n\nif torch.cuda.is_available():\n    vram_gb = torch.cuda.get_device_properties(0).total_mem / 1e9\n    gpu_name = torch.cuda.get_device_name(0)\n    print(f\"GPU: {gpu_name} ({vram_gb:.0f} GB)\")\nelse:\n    vram_gb = 0\n    print(\"No GPU detected — config will default to smallest profile\")\n\n# Auto-select profile based on VRAM, or override manually\nif vram_gb >= 40:  # H100, A100\n    PROFILE = \"h100\"\n    MODEL_NAME = \"unsloth/Qwen2.5-3B-Instruct\"\n    MAX_SEQ_LENGTH = 4096\n    MAX_PROMPT_LENGTH = 3584\n    MAX_COMPLETION_LENGTH = 512\n    BATCH_SIZE = 4\n    GRAD_ACCUM = 2\n    NUM_GENERATIONS = 8\nelif vram_gb >= 14:  # T4, L4\n    PROFILE = \"t4\"\n    MODEL_NAME = \"unsloth/Qwen2.5-0.5B-Instruct\"\n    MAX_SEQ_LENGTH = 2048\n    MAX_PROMPT_LENGTH = 1792\n    MAX_COMPLETION_LENGTH = 256\n    BATCH_SIZE = 2\n    GRAD_ACCUM = 4\n    NUM_GENERATIONS = 4\nelse:\n    PROFILE = \"cpu\"\n    MODEL_NAME = \"unsloth/Qwen2.5-0.5B-Instruct\"\n    MAX_SEQ_LENGTH = 2048\n    MAX_PROMPT_LENGTH = 1792\n    MAX_COMPLETION_LENGTH = 256\n    BATCH_SIZE = 1\n    GRAD_ACCUM = 8\n    NUM_GENERATIONS = 2\n\nprint(f\"Profile: {PROFILE} | Model: {MODEL_NAME} | Context: {MAX_SEQ_LENGTH}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "source": "import json\nimport re\nimport random\nfrom datasets import Dataset\n\n# Load episodes\nwith open(\"episodes.json\") as f:\n    episodes = json.load(f)\n\n# Check format — old format has 'messages'/'tasks', new format has 'decision_points'\nif episodes and \"decision_points\" not in episodes[0]:\n    old_keys = list(episodes[0].keys())\n    raise ValueError(\n        f\"episodes.json is in the old format (keys: {old_keys}).\\n\"\n        f\"Regenerate with: python generate_episodes.py -n 100\\n\"\n        f\"The old format used 'messages'/'tasks'/'schema_events'; \"\n        f\"the notebook requires 'decision_points' from generate_episodes.py.\"\n    )\n\n# Flatten to individual training prompts\nprompts = []\nfor ep in episodes:\n    for dp in ep[\"decision_points\"]:\n        prompts.append({\n            \"prompt\": dp[\"prompt\"],\n            \"hour\": dp[\"hour\"],\n            \"visible_count\": dp[\"visible_count\"],\n            \"episode_id\": ep[\"episode_id\"],\n            \"seed\": ep[\"seed\"],\n            \"drift_events\": ep[\"drift_events\"],\n            \"superseded\": ep.get(\"superseded_messages\", {}),\n            \"messages\": dp[\"visible_messages\"],\n        })\n\nif not prompts:\n    raise ValueError(\"No decision_points found in episodes; cannot train.\")\n\nprint(f\"Loaded {len(episodes)} episodes -> {len(prompts)} training prompts\")\nprint(f\"Average {len(prompts)/len(episodes):.1f} decision points per episode\")",
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Reward Function\n\nScores agent actions based on:\n- **Urgency base** (critical=10, high=5, medium=3, low=1)\n- **Deadline timing** (early=bonus, late=penalty)\n- **Drift adaptation** (+50% for handling policy-change messages)\n- **Stale info penalty** (-50% for acting on superseded messages)\n- **Response quality** (penalty for short/empty responses)",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "def score_action(completion: str, prompt_data: dict) -> float:\n    \"\"\"\n    Score a model completion against the inbox state.\n    \n    The model should output: respond_to_message(msg_id, \"response text\")\n    We parse the message_id and response, then score based on the reward function.\n    \"\"\"\n    messages = prompt_data[\"messages\"]\n    hour = prompt_data[\"hour\"]\n    superseded = prompt_data.get(\"superseded\", {})\n    \n    # Parse the model output for message_id\n    msg_id = None\n    response_text = \"\"\n    \n    # Try to parse respond_to_message(msg_id, response)\n    match = re.search(r'respond_to_message\\s*\\(\\s*[\"\\']?(msg_\\d+)[\"\\']?\\s*,\\s*[\"\\'](.+?)[\"\\']', completion, re.DOTALL)\n    if match:\n        msg_id = match.group(1)\n        response_text = match.group(2)\n    else:\n        # Try simpler format: just a message ID mentioned\n        id_match = re.search(r'(msg_\\d+)', completion)\n        if id_match:\n            msg_id = id_match.group(1)\n            # No explicit response text — penalize via quality check below\n            response_text = \"\"\n    \n    if not msg_id:\n        return -1.0  # couldn't parse any action\n    \n    # Find the message in the inbox\n    target_msg = None\n    for msg in messages:\n        if msg[\"id\"] == msg_id:\n            target_msg = msg\n            break\n    \n    if target_msg is None:\n        return -0.5  # referenced a message not in inbox\n    \n    # Base reward by urgency\n    urgency_rewards = {\"critical\": 10.0, \"high\": 5.0, \"medium\": 3.0, \"low\": 1.0}\n    reward = urgency_rewards.get(target_msg[\"urgency\"], 1.0)\n    \n    # Deadline timing\n    deadline = target_msg.get(\"deadline_hours\")\n    if deadline is not None:\n        if hour <= deadline:\n            time_remaining_frac = (deadline - hour) / max(deadline, 1.0)\n            reward *= 1.0 + 0.5 * time_remaining_frac\n        else:\n            reward *= 0.25  # late penalty\n    \n    # Response quality\n    if len(response_text.strip()) < 10:\n        reward *= 0.5\n    \n    # Drift adaptation bonus\n    if target_msg.get(\"drift_flag\"):\n        reward *= 1.5\n    \n    # Stale info penalty\n    if target_msg[\"id\"] in superseded:\n        reward *= 0.5\n    \n    # Penalize choosing low-urgency when unhandled critical messages exist\n    unhandled_critical = any(\n        m[\"urgency\"] == \"critical\" and not m.get(\"handled\") and not m.get(\"superseded\")\n        for m in messages\n    )\n    if unhandled_critical and target_msg[\"urgency\"] in (\"low\", \"medium\"):\n        reward *= 0.3\n    \n    return round(reward, 2)\n\n\n# Test the reward function\ntest_data = prompts[0]\nprint(\"Testing reward function on first decision point:\")\nprint(f\"  Hour: {test_data['hour']}, Messages: {test_data['visible_count']}\")\n\n# Simulate good action (pick critical message)\ncritical_msgs = [m for m in test_data[\"messages\"] if m[\"urgency\"] == \"critical\"]\nif critical_msgs:\n    good_action = f'respond_to_message(\"{critical_msgs[0][\"id\"]}\", \"Acknowledged. Evacuating immediately with documents and medication.\")'\n    good_score = score_action(good_action, test_data)\n    print(f\"  Good action (critical msg): {good_score:.2f} pts\")\n\n# Simulate bad action (pick low-urgency message)\nlow_msgs = [m for m in test_data[\"messages\"] if m[\"urgency\"] == \"low\"]\nif low_msgs:\n    bad_action = f'respond_to_message(\"{low_msgs[0][\"id\"]}\", \"ok\")'\n    bad_score = score_action(bad_action, test_data)\n    print(f\"  Bad action (low msg, short response): {bad_score:.2f} pts\")\n\n# Simulate unparseable action\njunk_score = score_action(\"I think we should do something\", test_data)\nprint(f\"  Unparseable action: {junk_score:.2f} pts\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Load Model & Configure GRPO",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "from unsloth import FastLanguageModel\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    model_name=MODEL_NAME,\n    max_seq_length=MAX_SEQ_LENGTH,\n    load_in_4bit=True,\n)\n\n# Add LoRA adapters — bigger r for bigger models\nlora_r = 32 if PROFILE == \"h100\" else 16\n\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=lora_r,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                     \"gate_proj\", \"up_proj\", \"down_proj\"],\n    lora_alpha=lora_r,\n    lora_dropout=0,\n    bias=\"none\",\n    use_gradient_checkpointing=\"unsloth\",\n)\nprint(f\"Model loaded: {MODEL_NAME} | LoRA r={lora_r} | ctx={MAX_SEQ_LENGTH}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": "# Build the training dataset\n# Each row needs a \"prompt\" field formatted as chat messages\ntrain_data = []\nfor p in prompts:\n    train_data.append({\n        \"prompt\": [\n            {\"role\": \"user\", \"content\": p[\"prompt\"]},\n        ],\n        # Store metadata for reward calculation (not used by trainer directly)\n        \"_hour\": p[\"hour\"],\n        \"_episode_id\": p[\"episode_id\"],\n    })\n\n# Shuffle and split\nrandom.seed(42)\nrandom.shuffle(train_data)\n\ndataset = Dataset.from_list(train_data)\nprint(f\"Training dataset: {len(dataset)} prompts\")\nprint(f\"Sample prompt length: {len(train_data[0]['prompt'][0]['content'])} chars\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## GRPO Training Loop\n\nThe reward function scores each completion by:\n1. Parsing which message the model chose to handle\n2. Checking urgency, deadline timing, drift flags\n3. Penalizing bad choices (low-urgency when critical exists, stale info)",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "from trl import GRPOConfig, GRPOTrainer\n\n# Build a lookup from (episode_id, hour) -> prompt metadata for reward scoring\nprompt_lookup = {}\nfor p in prompts:\n    key = (p[\"episode_id\"], p[\"hour\"])\n    prompt_lookup[key] = p\n\n\ndef reward_fn(prompts, completions, _episode_id, _hour, **kwargs):\n    \"\"\"\n    GRPO reward function. Scores each completion against its inbox state.\n\n    TRL passes extra dataset columns as keyword arguments, so _episode_id and\n    _hour come directly from the dataset — no need to reverse-lookup from text.\n    \"\"\"\n    rewards = []\n    for completion, ep_id, hour in zip(completions, _episode_id, _hour):\n        key = (ep_id, hour)\n        prompt_data = prompt_lookup.get(key)\n\n        if prompt_data is None:\n            rewards.append(0.0)\n            continue\n\n        if isinstance(completion, list):\n            comp_text = completion[-1][\"content\"] if completion else \"\"\n        else:\n            comp_text = str(completion)\n\n        score = score_action(comp_text, prompt_data)\n        rewards.append(score)\n\n    return rewards\n\n\nprint(f\"Prompt lookup: {len(prompt_lookup)} unique keys (expect {len(prompts)})\")\n\n# GRPO training config — all values from GPU profile\ntraining_args = GRPOConfig(\n    output_dir=\"crisisinbox-grpo-output\",\n    num_train_epochs=3,\n    per_device_train_batch_size=BATCH_SIZE,\n    gradient_accumulation_steps=GRAD_ACCUM,\n    learning_rate=5e-6,\n    max_completion_length=MAX_COMPLETION_LENGTH,\n    max_prompt_length=MAX_PROMPT_LENGTH,\n    num_generations=NUM_GENERATIONS,\n    logging_steps=10,\n    save_steps=100,\n    report_to=\"none\",\n    bf16=True,\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_fn,\n    args=training_args,\n    train_dataset=dataset,\n)\n\nprint(f\"Trainer configured — batch={BATCH_SIZE}, gen={NUM_GENERATIONS}, prompt≤{MAX_PROMPT_LENGTH}tok\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": "# Train!\ntrainer.train()\nprint(\"Training complete\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "## Evaluate Trained Model\n\nSample prompts and check whether the model picks high-urgency messages and produces well-formatted actions.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "# Evaluate on a few test prompts\nFastLanguageModel.for_inference(model)\n\neval_prompts = random.sample(prompts, min(10, len(prompts)))\ntotal_score = 0\n\nprint(f\"=== Trained Model Evaluation ({MODEL_NAME}) ===\\n\")\nfor p in eval_prompts:\n    messages = [{\"role\": \"user\", \"content\": p[\"prompt\"]}]\n    inputs = tokenizer.apply_chat_template(messages, return_tensors=\"pt\", add_generation_prompt=True).to(\"cuda\")\n\n    with torch.no_grad():\n        output = model.generate(inputs, max_new_tokens=MAX_COMPLETION_LENGTH, temperature=0.7, do_sample=True)\n\n    completion = tokenizer.decode(output[0][inputs.shape[1]:], skip_special_tokens=True)\n    score = score_action(completion, p)\n    total_score += score\n\n    # Show a summary\n    msg_match = re.search(r'(msg_\\d+)', completion)\n    chosen_id = msg_match.group(1) if msg_match else \"none\"\n    chosen_msg = next((m for m in p[\"messages\"] if m[\"id\"] == chosen_id), None)\n    urgency = chosen_msg[\"urgency\"] if chosen_msg else \"?\"\n\n    print(f\"Hour {p['hour']:5.1f} | Chose: {chosen_id} ({urgency:8s}) | Score: {score:+.1f}\")\n\nprint(f\"\\nAverage score: {total_score / len(eval_prompts):.2f}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": "# Save the trained model\nmodel.save_pretrained(\"crisisinbox-grpo-trained\")\ntokenizer.save_pretrained(\"crisisinbox-grpo-trained\")\nprint(\"Model saved to crisisinbox-grpo-trained/\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "crisis-inbox"
+version = "0.1.0"
+description = "RL environment for training LLMs to manage personal task overload during natural disasters"
+requires-python = ">=3.10"
+dependencies = [
+    "openenv-core[core]>=0.2.1",
+    "fastapi>=0.115.0",
+    "pydantic>=2.0.0",
+    "uvicorn>=0.24.0",
+    "requests>=2.31.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+server = "crisis_inbox.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["crisis_inbox", "crisis_inbox.server"]
+package-dir = { "crisis_inbox" = ".", "crisis_inbox.server" = "server" }

server/Dockerfile

@@ -2,21 +2,19 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
-# Install uv for fast dependency management
-RUN pip install uv
+COPY requirements.txt /app/
+RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy project files
-COPY pyproject.toml /app/
-COPY crisis_inbox/ /app/crisis_inbox/
+COPY models.py /app/
+COPY messages.py /app/
+COPY drift_events.py /app/
+COPY client.py /app/
+COPY __init__.py /app/
+COPY server/ /app/server/
 
-# Install dependencies
-RUN uv pip install --system -e .
-
-# Expose port
 EXPOSE 8000
 
-# Health check
 HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
     CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1
 
-CMD ["uvicorn", "crisis_inbox.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
+CMD ["uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "8000"]

server/crisis_inbox_environment.py

@@ -293,7 +293,19 @@ class CrisisInboxEnvironment(MCPEnvironment):
         for msg in ALL_MESSAGES:
             if msg.id in all_drift_ids and msg.id not in drift_msg_ids:
                 continue  # Skip drift messages from events not selected this episode
-            self._all_messages.append(msg.model_copy())
+            m = msg.model_copy()
+            # Add jitter to arrival times and deadlines for episode variation
+            # Keep hour-0 messages at 0, jitter others by +/- 15% (clamped to 0-48)
+            if m.timestamp_hours > 0:
+                jitter = self._rng.uniform(-0.15, 0.15) * m.timestamp_hours
+                m.timestamp_hours = max(0.1, min(47.5, m.timestamp_hours + jitter))
+            if m.deadline_hours is not None and m.deadline_hours > 0:
+                d_jitter = self._rng.uniform(-0.1, 0.1) * m.deadline_hours
+                m.deadline_hours = max(
+                    m.timestamp_hours + 0.5,
+                    min(72.0, m.deadline_hours + d_jitter),
+                )
+            self._all_messages.append(m)
 
         self._visible_messages = []
         self._handled = {}

training/crisisinbox_training.py

@@ -0,0 +1,416 @@
+"""
+CrisisInbox GRPO Training Script
+Person B: ML Pipeline
+
+Run this in Google Colab:
+1. Upload this file
+2. Upload episodes.json from repo
+3. Run: python crisisinbox_training.py
+"""
+
+import torch
+import re
+import json
+import numpy as np
+from datasets import Dataset
+from unsloth import FastLanguageModel
+from trl import GRPOConfig, GRPOTrainer
+
+# Download episodes from GitHub repo
+print("Loading episodes...")
+import urllib.request
+urllib.request.urlretrieve(
+    "https://raw.githubusercontent.com/eptan/crisis-inbox/main/episodes.json",
+    "episodes.json"
+)
+
+with open("episodes.json", "r") as f:
+    EPISODES = json.load(f)
+
+print(f"✓ Loaded {len(EPISODES)} episodes")
+
+# =============================================================================
+# PROMPT BUILDING
+# =============================================================================
+
+CRISIS_SYSTEM_PROMPT = """
+You are an assistant helping a working parent during a wildfire.
+You must triage messages, act on safety-critical items first,
+respect deadlines, adapt to policy changes, and write appropriate
+tones to different people.
+
+Respond ONLY in this format:
+<plan>
+1. [time=min X] ACTION_DESCRIPTION
+2. [time=min Y] ACTION_DESCRIPTION
+...
+</plan>
+"""
+
+def build_crisis_prompt(episode):
+    """Build prompt from episode data."""
+    msgs_str = []
+    for m in episode.get("messages", []):
+        deadline_info = f" (DEADLINE: {m['deadline']}h)" if m.get("deadline") else ""
+        urgency = "🔴 URGENT" if m.get("type") == "safety" else ""
+        msgs_str.append(
+            f"[t={m['time']}h] {urgency} From {m['sender']} via {m['channel']}: {m['content']}{deadline_info}"
+        )
+    
+    drift_str = []
+    for d in episode.get("schema_events", []):
+        drift_str.append(f"[t={d['time']}h] POLICY UPDATE: {d['kind']} -> {d.get('new_value', 'changed')}")
+    
+    user_content = (
+        "Here is your 48-hour message history:\n\n"
+        + "\n".join(msgs_str)
+        + "\n\nPolicy changes observed:\n"
+        + ("\n".join(drift_str) if drift_str else "None yet.")
+        + "\n\nDecide what to do in order. Remember the required <plan> format."
+    )
+    
+    return [
+        {"role": "system", "content": CRISIS_SYSTEM_PROMPT},
+        {"role": "user", "content": user_content},
+    ]
+
+# =============================================================================
+# PLAN PARSER
+# =============================================================================
+
+def parse_plan(model_output):
+    """Parse <plan> tag output into list of action dicts."""
+    actions = []
+    
+    plan_match = re.search(r'<plan>(.*?)</plan>', model_output, re.DOTALL | re.IGNORECASE)
+    if not plan_match:
+        return []
+    
+    plan_content = plan_match.group(1).strip()
+    
+    lines = plan_content.split('\n')
+    for line in lines:
+        line = line.strip()
+        if not line or not line[0].isdigit():
+            continue
+        
+        # Extract time: [time=min X]
+        time_match = re.search(r'time=min (\d+)', line)
+        time_min = int(time_match.group(1)) if time_match else 0
+        
+        # Extract action description
+        action_desc = line.split(']', 1)[1].strip() if ']' in line else line
+        
+        # Infer target
+        target = "unknown"
+        common_targets = ["spouse", "boss", "mom", "dad", "sister", "neighbor", "FEMA", "Insurance", "Airline", "School"]
+        for t in common_targets:
+            if t.lower() in action_desc.lower():
+                target = t
+                break
+        
+        actions.append({
+            "time_min": time_min,
+            "action": action_desc,
+            "target": target,
+            "task_id": None
+        })
+    
+    return actions
+
+# =============================================================================
+# 5-COMPONENT REWARD FUNCTION
+# =============================================================================
+
+W_SAFETY   = 1.5
+W_DEADLINE = 1.5
+W_DRIFT    = 1.0
+W_TONE     = 0.5
+W_COVER    = 1.0
+
+def score_safety_priority(episode, parsed_actions):
+    """+10 if safety in first 3 actions, -10 if ignored."""
+    safety_tasks = [t_id for t_id, t in episode.get("tasks", {}).items() 
+                    if t.get("category") == "safety"]
+    
+    acted_tasks = set()
+    for a in parsed_actions:
+        for t_id, t_info in episode.get("tasks", {}).items():
+            if a.get("target", "").lower() in t_info.get("sender", "").lower():
+                a["task_id"] = t_id
+                acted_tasks.add(t_id)
+                break
+    
+    score = 0.0
+    first_three_targets = [a.get("target", "").lower() for a in parsed_actions[:3]]
+    safety_targets = [episode["tasks"][t]["sender"].lower() for t in safety_tasks]
+    
+    if any(t in first_three_targets for t in safety_targets):
+        score += 10.0
+    
+    safety_acted = acted_tasks.intersection(set(safety_tasks))
+    if not safety_acted and safety_tasks:
+        score -= 10.0
+    
+    return score
+
+def score_deadlines(episode, parsed_actions):
+    """+5 before deadline, -5 after/missed."""
+    score = 0.0
+    completion_time = {}
+    
+    for a in parsed_actions:
+        for t_id, t_info in episode.get("tasks", {}).items():
+            if a.get("target", "").lower() in t_info.get("sender", "").lower():
+                if t_id not in completion_time:
+                    completion_time[t_id] = a["time_min"] / 60.0
+                break
+    
+    for t_id, deadline in episode.get("deadlines", {}).items():
+        if t_id in completion_time:
+            if completion_time[t_id] <= deadline:
+                score += 5.0
+            else:
+                score -= 5.0
+        else:
+            score -= 5.0
+    
+    return score
+
+def score_schema_drift(episode, parsed_actions):
+    """+10 for adapting to policy changes."""
+    score = 0.0
+    
+    for event in episode.get("schema_events", []):
+        event_time = event["time"]
+        event_kind = event.get("kind", "")
+        
+        actions_after = [a for a in parsed_actions if a["time_min"] / 60.0 > event_time]
+        
+        if not actions_after:
+            continue
+        
+        if "insurance_deadline_change" in event_kind:
+            new_deadline = event.get("new_value", 72)
+            old_deadline = event.get("old_value", 72)
+            
+            insurance_actions = [a for a in actions_after 
+                               if "insurance" in a.get("action", "").lower()]
+            
+            for a in insurance_actions:
+                action_hour = a["time_min"] / 60.0
+                if action_hour <= new_deadline:
+                    score += 10.0
+                elif action_hour > old_deadline:
+                    score -= 10.0
+                else:
+                    score -= 5.0
+    
+    return score
+
+def score_tone(episode, model_raw_output):
+    """Simple tone heuristics."""
+    score = 0.0
+    text = model_raw_output.lower()
+    
+    if "dear" in text or "sincerely" in text or "regards" in text:
+        score += 1.0
+    
+    if "love you" in text or "i'm so sorry" in text or "worried" in text:
+        score += 1.0
+    
+    if "confirm" in text or "verified" in text or "documentation" in text:
+        score += 1.0
+    
+    return score
+
+def score_coverage(episode, parsed_actions):
+    """+2 per task touched, -1 per ignored."""
+    described_tasks = set(episode.get("tasks", {}).keys())
+    
+    acted_tasks = set()
+    for a in parsed_actions:
+        for t_id, t_info in episode.get("tasks", {}).items():
+            if a.get("target", "").lower() in t_info.get("sender", "").lower():
+                acted_tasks.add(t_id)
+                break
+    
+    score = 0.0
+    for t_id in described_tasks:
+        if t_id in acted_tasks:
+            score += 2.0
+        else:
+            score -= 1.0
+    
+    return score
+
+def total_reward(episode, model_raw_output):
+    """Calculate total reward from 5 components."""
+    parsed_actions = parse_plan(model_raw_output)
+    
+    if not parsed_actions:
+        return -20.0
+    
+    r_safety   = score_safety_priority(episode, parsed_actions)
+    r_deadline = score_deadlines(episode, parsed_actions)
+    r_drift    = score_schema_drift(episode, parsed_actions)
+    r_tone     = score_tone(episode, model_raw_output)
+    r_cover    = score_coverage(episode, parsed_actions)
+    
+    total = (
+        W_SAFETY   * r_safety +
+        W_DEADLINE * r_deadline +
+        W_DRIFT    * r_drift +
+        W_TONE     * r_tone +
+        W_COVER    * r_cover
+    )
+    
+    return total
+
+# =============================================================================
+# GRPO TRAINING SETUP
+# =============================================================================
+
+def build_dataset(episodes):
+    rows = []
+    for idx, ep in enumerate(episodes):
+        prompt = build_crisis_prompt(ep)
+        rows.append({
+            "id": idx,
+            "prompt": prompt,
+            "episode": ep,
+        })
+    return Dataset.from_list(rows)
+
+print("Building dataset...")
+crisis_dataset = build_dataset(EPISODES)
+print(f"✓ Built dataset with {len(crisis_dataset)} episodes")
+
+# Load model
+print("Loading model...")
+MODEL_NAME = "unsloth/Qwen2.5-0.5B-Instruct"
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=MODEL_NAME,
+    max_seq_length=2048,
+    dtype=None,
+    load_in_4bit=True,
+)
+
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=16,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                    "gate_proj", "up_proj", "down_proj"],
+    lora_alpha=16,
+    lora_dropout=0,
+    bias="none",
+    use_gradient_checkpointing="unsloth",
+    random_state=3407,
+)
+
+print(f"✓ Loaded model: {MODEL_NAME}")
+
+# Reward function
+EPISODES_LIST = [row["episode"] for row in crisis_dataset]
+
+def crisis_reward_fn(prompts, completions, **kwargs):
+    rewards = []
+    
+    for prompt, completion in zip(prompts, completions):
+        # Convert completion to string - completion is list of token IDs
+        if isinstance(completion, list):
+            # Access tokenizer from global scope
+            from unsloth import FastLanguageModel
+            model_obj = FastLanguageModel.get_model()
+            tokenizer_obj = model_obj.tokenizer if hasattr(model_obj, 'tokenizer') else None
+            if tokenizer_obj is not None:
+                completion_str = tokenizer_obj.decode(completion, skip_special_tokens=True)
+            else:
+                completion_str = str(completion)
+        else:
+            completion_str = str(completion)
+        
+        episode = None
+        for ep in EPISODES_LIST:
+            test_prompt = build_crisis_prompt(ep)
+            if test_prompt[1]["content"] == prompt[1]["content"]:
+                episode = ep
+                break
+        
+        if episode is None:
+            idx = kwargs.get("episode_idx", 0)
+            episode = EPISODES_LIST[idx % len(EPISODES_LIST)]
+        
+        reward = total_reward(episode, completion_str)
+        rewards.append(float(reward))
+    
+    return rewards
+
+# Training config
+training_args = GRPOConfig(
+    use_vllm=False,  # Disabled for Colab compatibility
+    num_train_epochs=3,
+    max_steps=200,  # Start with 200 for test
+    per_device_train_batch_size=2,
+    gradient_accumulation_steps=4,
+    num_generations=4,
+    max_completion_length=512,
+    temperature=0.7,
+    learning_rate=1e-5,
+    logging_steps=10,
+    save_steps=50,
+    output_dir="./crisisinbox_grpo",
+    report_to="none",
+)
+
+trainer = GRPOTrainer(
+    model=model,
+    processing_class=tokenizer,
+    reward_funcs=crisis_reward_fn,
+    train_dataset=crisis_dataset,
+    args=training_args,
+)
+
+print("✓ GRPO Trainer ready")
+
+# =============================================================================
+# TRAINING RUN
+# =============================================================================
+
+print("🚀 Starting CrisisInbox GRPO Training...")
+print("This will take 1-4 hours depending on max_steps\n")
+
+trainer.train()
+
+print("\n✅ Training completed!")
+
+# =============================================================================
+# SAVE RESULTS
+# =============================================================================
+
+model.save_pretrained("./crisisinbox_model_final")
+tokenizer.save_pretrained("./crisisinbox_model_final")
+print("✓ Model saved")
+
+# Generate demo examples
+DEMO_EXAMPLES = []
+for ep in EPISODES[:3]:
+    prompt = build_crisis_prompt(ep)
+    text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=200, temperature=0.7)
+    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    
+    reward = total_reward(ep, response)
+    DEMO_EXAMPLES.append({
+        "episode_id": ep["episode_id"],
+        "response": response,
+        "reward": reward,
+        "parsed_actions": parse_plan(response)
+    })
+
+with open("demo_examples_trained.json", "w") as f:
+    json.dump(DEMO_EXAMPLES, f, indent=2)
+
+print(f"✓ Saved {len(DEMO_EXAMPLES)} demo examples")