Upload folder using huggingface_hub

__init__.py

@@ -8,9 +8,11 @@ from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction
 try:
     from .client import CrisisInboxEnv
     from .models import Channel, Message, Urgency
+    from .rewards import calculate_reward, tone_multiplier
 except ImportError:
     from client import CrisisInboxEnv
     from models import Channel, Message, Urgency
+    from rewards import calculate_reward, tone_multiplier
 
 __all__ = [
     "CrisisInboxEnv",
@@ -19,4 +21,6 @@ __all__ = [
     "Message",
     "Channel",
     "Urgency",
+    "calculate_reward",
+    "tone_multiplier",
 ]

notebooks/crisisinbox_grpo_connected.ipynb

@@ -3,13 +3,21 @@
   {
    "cell_type": "markdown",
    "id": "ym4tunggrm",
-   "source": "# CrisisInbox GRPO Training (Connected to HF Space)\n\nTrain a small LLM (Qwen2.5-0.5B) to triage crisis inbox messages using Group Relative Policy Optimization.\n\n**This notebook connects to the live CrisisInbox environment** deployed on HuggingFace Spaces at `https://eptan-crisis-inbox.hf.space` to collect training episodes in real-time, then trains the model using GRPO.\n\n**Stack:** HF TRL + PEFT (LoRA on full bf16 model — no quantization needed for 0.5B)\n\n**What this does:**\n1. Connects to the deployed CrisisInbox environment via WebSocket\n2. Collects episodes by interacting with the environment (reset, list tools, call tools)\n3. Builds training prompts from live environment observations\n4. Trains the model with GRPO using a reward function\n5. Evaluates the trained model against the live environment\n\nOpen in Google Colab or Northflank with a GPU runtime.",
+   "source": "# CrisisInbox GRPO Training (Connected to HF Space)\n\nTrain a small LLM (Qwen2.5-0.5B) to triage crisis inbox messages using Group Relative Policy Optimization.\n\n**This notebook connects to the live CrisisInbox environment** deployed on HuggingFace Spaces at `https://eptan-crisis-inbox.hf.space` to collect training episodes in real-time, then trains the model using GRPO.\n\n**Stack:** HF TRL + PEFT (LoRA on full bf16 model \u2014 no quantization needed for 0.5B)\n\n**What this does:**\n1. Connects to the deployed CrisisInbox environment via WebSocket\n2. Collects episodes by interacting with the environment (reset, list tools, call tools)\n3. Builds training prompts from live environment observations\n4. Trains the model with GRPO using a reward function\n5. Evaluates the trained model against the live environment\n\nOpen in Google Colab or Northflank with a GPU runtime.",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "p0j1w7pr7ib",
-   "source": "# Install dependencies (pure HF TRL + PEFT, no quantization needed for 0.5B model)\n!pip install trl transformers datasets accelerate peft -q\n!pip install \"openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git\" -q\n!pip install huggingface_hub matplotlib -q\nprint(\"Setup complete\")",
+   "source": [
+    "# Install dependencies (pure HF TRL + PEFT, no quantization needed for 0.5B model)\n",
+    "!pip install trl transformers datasets accelerate peft -q\n",
+    "!pip install \"openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git\" -q\n",
+    "!pip install huggingface_hub matplotlib -q\n",
+    "# Install crisis-inbox package so we can import the shared reward function\n",
+    "!pip install \"crisis-inbox @ git+https://github.com/eptan/crisis-inbox.git\" -q\n",
+    "print(\"Setup complete\")\n"
+   ],
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -39,7 +47,7 @@
   {
    "cell_type": "code",
    "id": "pmbt9gcp9hb",
-   "source": "import json\nimport time as _time\nfrom openenv.core.mcp_client import MCPToolClient\n\nBASE_URL = \"https://eptan-crisis-inbox.hf.space\"\n\n# Wake up the HF Space (may be sleeping) and verify connectivity\nprint(\"Connecting to HF Space (may take a moment if cold-starting)...\")\nfor attempt in range(3):\n    try:\n        with MCPToolClient(base_url=BASE_URL, connect_timeout_s=60.0).sync() as env:\n            env.reset(seed=0)\n            tools = env.list_tools()\n            print(f\"Connected! Available tools: {[t.name for t in tools]}\")\n            for t in tools:\n                print(f\"  - {t.name}: {t.description[:80]}...\")\n\n            status = json.loads(env.call_tool(\"get_status\"))\n            print(f\"\\nEnvironment ready — {status['messages_total_arrived']} messages at hour {status['current_hour']}\")\n        break\n    except Exception as e:\n        if attempt < 2:\n            print(f\"  Attempt {attempt + 1} failed ({e}), retrying in 10s...\")\n            _time.sleep(10)\n        else:\n            raise RuntimeError(f\"Could not connect to {BASE_URL} after 3 attempts: {e}\")",
+   "source": "import json\nimport time as _time\nfrom openenv.core.mcp_client import MCPToolClient\n\nBASE_URL = \"https://eptan-crisis-inbox.hf.space\"\n\n# Wake up the HF Space (may be sleeping) and verify connectivity\nprint(\"Connecting to HF Space (may take a moment if cold-starting)...\")\nfor attempt in range(3):\n    try:\n        with MCPToolClient(base_url=BASE_URL, connect_timeout_s=60.0).sync() as env:\n            env.reset(seed=0)\n            tools = env.list_tools()\n            print(f\"Connected! Available tools: {[t.name for t in tools]}\")\n            for t in tools:\n                print(f\"  - {t.name}: {t.description[:80]}...\")\n\n            status = json.loads(env.call_tool(\"get_status\"))\n            print(f\"\\nEnvironment ready \u2014 {status['messages_total_arrived']} messages at hour {status['current_hour']}\")\n        break\n    except Exception as e:\n        if attempt < 2:\n            print(f\"  Attempt {attempt + 1} failed ({e}), retrying in 10s...\")\n            _time.sleep(10)\n        else:\n            raise RuntimeError(f\"Could not connect to {BASE_URL} after 3 attempts: {e}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -69,7 +77,113 @@
   {
    "cell_type": "code",
    "id": "2xd2afp4g99",
-   "source": "import re\n\ndef score_action(completion: str, prompt_data: dict) -> float:\n    \"\"\"Score a model completion against the inbox state.\n\n    This mirrors the server-side _calculate_reward() in\n    crisis_inbox_environment.py for offline GRPO training.\n    The server is the single source of truth — keep these in sync.\n\n    The model should output: respond_to_message(msg_id, \"response text\")\n    We parse the message_id and response, then score using the same\n    reward signals as the server: urgency, deadline, drift, stale, priority.\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    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        id_match = re.search(r'(msg_\\d+)', completion)\n        if id_match:\n            msg_id = id_match.group(1)\n            response_text = completion\n\n    if not msg_id:\n        return -1.0\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\n\n    # --- Reward signals (mirroring server _calculate_reward) ---\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\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    # Priority penalty: choosing low/medium when unhandled critical exists\n    unhandled = [m for m in messages if not m.get(\"handled\", False) and m[\"id\"] != msg_id]\n    has_unhandled_critical = any(m[\"urgency\"] == \"critical\" for m in unhandled)\n    if has_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\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\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\njunk_score = score_action(\"I think we should do something\", test_data)\nprint(f\"  Unparseable action: {junk_score:.2f} pts\")",
+   "source": [
+    "import re\n",
+    "from crisis_inbox.rewards import calculate_reward\n",
+    "from crisis_inbox.models import Message, Urgency, Channel\n",
+    "\n",
+    "\n",
+    "def score_action(completion: str, prompt_data: dict) -> float:\n",
+    "    \"\"\"Score a model completion against the inbox state.\n",
+    "\n",
+    "    Parses the model output for respond_to_message(msg_id, response),\n",
+    "    constructs a Message object, and delegates to the shared\n",
+    "    calculate_reward() from rewards.py (single source of truth).\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 and response text\n",
+    "    msg_id = None\n",
+    "    response_text = \"\"\n",
+    "\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",
+    "        id_match = re.search(r'(msg_\\d+)', completion)\n",
+    "        if id_match:\n",
+    "            msg_id = id_match.group(1)\n",
+    "            response_text = completion\n",
+    "\n",
+    "    if not msg_id:\n",
+    "        return -1.0\n",
+    "\n",
+    "    # Find the message dict in the inbox\n",
+    "    target_dict = None\n",
+    "    for msg in messages:\n",
+    "        if msg[\"id\"] == msg_id:\n",
+    "            target_dict = msg\n",
+    "            break\n",
+    "\n",
+    "    if target_dict is None:\n",
+    "        return -0.5\n",
+    "\n",
+    "    # Construct a Message object for the shared reward function\n",
+    "    target_msg = Message(\n",
+    "        id=target_dict[\"id\"],\n",
+    "        sender=target_dict[\"sender\"],\n",
+    "        channel=target_dict.get(\"channel\", \"email\"),\n",
+    "        subject=target_dict.get(\"subject\", \"\"),\n",
+    "        content=target_dict.get(\"content\", \"\"),\n",
+    "        urgency=target_dict[\"urgency\"],\n",
+    "        timestamp_hours=target_dict.get(\"timestamp_hours\", 0.0),\n",
+    "        deadline_hours=target_dict.get(\"deadline_hours\"),\n",
+    "        dependencies=target_dict.get(\"dependencies\", []),\n",
+    "        drift_flag=target_dict.get(\"drift_flag\", False),\n",
+    "        supersedes=target_dict.get(\"supersedes\"),\n",
+    "    )\n",
+    "\n",
+    "    # Build visible_messages as Message objects for priority penalty\n",
+    "    handled_ids = {m[\"id\"]: \"\" for m in messages if m.get(\"handled\", False)}\n",
+    "    visible = []\n",
+    "    for m in messages:\n",
+    "        visible.append(Message(\n",
+    "            id=m[\"id\"],\n",
+    "            sender=m[\"sender\"],\n",
+    "            channel=m.get(\"channel\", \"email\"),\n",
+    "            subject=m.get(\"subject\", \"\"),\n",
+    "            content=m.get(\"content\", \"\"),\n",
+    "            urgency=m[\"urgency\"],\n",
+    "            timestamp_hours=m.get(\"timestamp_hours\", 0.0),\n",
+    "            deadline_hours=m.get(\"deadline_hours\"),\n",
+    "            dependencies=m.get(\"dependencies\", []),\n",
+    "            drift_flag=m.get(\"drift_flag\", False),\n",
+    "            supersedes=m.get(\"supersedes\"),\n",
+    "        ))\n",
+    "\n",
+    "    return calculate_reward(\n",
+    "        msg=target_msg,\n",
+    "        current_hour=hour,\n",
+    "        response=response_text,\n",
+    "        superseded=superseded,\n",
+    "        visible_messages=visible,\n",
+    "        handled=handled_ids,\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "# Test the reward function\n",
+    "test_data = prompts[0]\n",
+    "print(\"Testing reward function on first decision point:\")\n",
+    "print(f\"  Hour: {test_data[\"hour\"]}, Messages: {test_data[\"visible_count\"]}\")\n",
+    "\n",
+    "critical_msgs = [m for m in test_data[\"messages\"] if m[\"urgency\"] == \"critical\"]\n",
+    "if 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",
+    "low_msgs = [m for m in test_data[\"messages\"] if m[\"urgency\"] == \"low\"]\n",
+    "if 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",
+    "junk_score = score_action(\"I think we should do something\", test_data)\n",
+    "print(f\"  Unparseable action: {junk_score:.2f} pts\")\n"
+   ],
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -83,7 +197,7 @@
   {
    "cell_type": "code",
    "id": "zey499u5w1a",
-   "source": "from transformers import AutoModelForCausalLM, AutoTokenizer\nfrom peft import LoraConfig\nimport torch\n\n# Auto-detect precision\n_use_bf16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False\n_compute_dtype = torch.bfloat16 if _use_bf16 else torch.float16\n\n# Load in full bf16/fp16 — no 4-bit quantization.\n# Qwen2.5-0.5B is ~1GB in bf16, fits easily on any GPU.\n# This avoids all bitsandbytes dtype mismatch issues with lm_head.\nmodel = AutoModelForCausalLM.from_pretrained(\n    \"Qwen/Qwen2.5-0.5B-Instruct\",\n    device_map=\"auto\",\n    torch_dtype=_compute_dtype,\n)\ntokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-0.5B-Instruct\")\n\n# Fix: TRL GRPOTrainer expects warnings_issued but newer transformers removed it.\nif not hasattr(model, \"warnings_issued\"):\n    model.warnings_issued = {}\n\n# GRPO requires left padding so completions align across the batch\ntokenizer.padding_side = \"left\"\nif tokenizer.pad_token_id is None:\n    tokenizer.pad_token = tokenizer.eos_token\n    tokenizer.pad_token_id = tokenizer.eos_token_id\n\n# LoRA config — passed to GRPOTrainer, not applied here\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=16,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                     \"gate_proj\", \"up_proj\", \"down_proj\"],\n    lora_dropout=0.0,\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\n\nprint(f\"Model loaded in {_compute_dtype} (no quantization)\")\nprint(f\"Precision: {'bf16' if _use_bf16 else 'fp16'}\")\nprint(f\"Model size: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params\")",
+   "source": "from transformers import AutoModelForCausalLM, AutoTokenizer\nfrom peft import LoraConfig\nimport torch\n\n# Auto-detect precision\n_use_bf16 = torch.cuda.is_bf16_supported() if torch.cuda.is_available() else False\n_compute_dtype = torch.bfloat16 if _use_bf16 else torch.float16\n\n# Load in full bf16/fp16 \u2014 no 4-bit quantization.\n# Qwen2.5-0.5B is ~1GB in bf16, fits easily on any GPU.\n# This avoids all bitsandbytes dtype mismatch issues with lm_head.\nmodel = AutoModelForCausalLM.from_pretrained(\n    \"Qwen/Qwen2.5-0.5B-Instruct\",\n    device_map=\"auto\",\n    torch_dtype=_compute_dtype,\n)\ntokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-0.5B-Instruct\")\n\n# Fix: TRL GRPOTrainer expects warnings_issued but newer transformers removed it.\nif not hasattr(model, \"warnings_issued\"):\n    model.warnings_issued = {}\n\n# GRPO requires left padding so completions align across the batch\ntokenizer.padding_side = \"left\"\nif tokenizer.pad_token_id is None:\n    tokenizer.pad_token = tokenizer.eos_token\n    tokenizer.pad_token_id = tokenizer.eos_token_id\n\n# LoRA config \u2014 passed to GRPOTrainer, not applied here\nlora_config = LoraConfig(\n    r=16,\n    lora_alpha=16,\n    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n                     \"gate_proj\", \"up_proj\", \"down_proj\"],\n    lora_dropout=0.0,\n    bias=\"none\",\n    task_type=\"CAUSAL_LM\",\n)\n\nprint(f\"Model loaded in {_compute_dtype} (no quantization)\")\nprint(f\"Precision: {'bf16' if _use_bf16 else 'fp16'}\")\nprint(f\"Model size: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -97,7 +211,7 @@
   {
    "cell_type": "code",
    "id": "6r2zhgg94fk",
-   "source": "# --- Pre-training baseline evaluation against live environment ---\nfrom openenv.core.env_server.mcp_types import CallToolAction\n\ndef generate_action_baseline(model, tokenizer, prompt_text):\n    \"\"\"Generate an action from the model (used for both baseline and trained eval).\"\"\"\n    msgs = [{\"role\": \"user\", \"content\": prompt_text}]\n    input_ids = tokenizer.apply_chat_template(msgs, return_tensors=\"pt\", add_generation_prompt=True)\n    if not isinstance(input_ids, torch.Tensor):\n        input_ids = input_ids[\"input_ids\"]\n    input_ids = input_ids.to(\"cuda\")\n    prompt_len = input_ids.shape[1]\n    with torch.no_grad():\n        output = model.generate(input_ids=input_ids, max_new_tokens=200, temperature=0.7,\n                                pad_token_id=tokenizer.pad_token_id, do_sample=True)\n    return tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)\n\n\ndef _extract_tool_result(obs):\n    \"\"\"Extract the JSON result dict from a CallToolObservation.\n\n    obs.result may be a plain string, a dict with 'data' key (FastMCP\n    CallToolResult serialization), or an object with a .data attribute.\n    Normalize to a dict.\n    \"\"\"\n    raw = getattr(obs, \"result\", None)\n\n    # FastMCP CallToolResult object → unwrap .data\n    if hasattr(raw, \"data\"):\n        raw = raw.data\n\n    # Serialized as {\"data\": \"...\"} dict\n    if isinstance(raw, dict) and \"data\" in raw:\n        raw = raw[\"data\"]\n\n    # Now raw should be a JSON string\n    if isinstance(raw, str):\n        try:\n            return json.loads(raw)\n        except (json.JSONDecodeError, TypeError):\n            return {}\n\n    # Already a dict\n    if isinstance(raw, dict):\n        return raw\n\n    return {}\n\n\ndef evaluate_on_live_env(model, tokenizer, base_url, seed, max_steps=20):\n    \"\"\"Run model against the live environment using OpenEnv's step() flow.\n\n    Uses env.step(CallToolAction(...)) for respond_to_message so that\n    rewards flow through Observation.reward — the standard OpenEnv contract.\n    Read-only tools (get_inbox, get_status, advance_time) use call_tool().\n    \"\"\"\n    with MCPToolClient(base_url=base_url, connect_timeout_s=60.0, message_timeout_s=120.0).sync() as env:\n        env.reset(seed=seed)\n        total_reward = 0.0\n        actions_taken = []\n\n        for step in range(max_steps):\n            inbox = json.loads(env.call_tool(\"get_inbox\"))\n            status = json.loads(env.call_tool(\"get_status\"))\n            current_hour = status[\"current_hour\"]\n\n            if status.get(\"done\"):\n                break\n\n            unhandled = [m for m in inbox if not m.get(\"handled\", False)]\n            if not unhandled:\n                env.call_tool(\"advance_time\", hours=2.0)\n                continue\n\n            completion = generate_action_baseline(model, tokenizer, build_prompt(inbox, current_hour))\n\n            match = re.search(r'respond_to_message\\s*\\(\\s*[\"\\']?(msg_\\d+)[\"\\']?\\s*,\\s*[\"\\'](.+?)[\"\\']', completion, re.DOTALL)\n            if not match:\n                id_match = re.search(r'(msg_\\d+)', completion)\n                if id_match:\n                    msg_id, response_text = id_match.group(1), completion[:200]\n                else:\n                    env.call_tool(\"advance_time\", hours=1.0)\n                    continue\n            else:\n                msg_id, response_text = match.group(1), match.group(2)\n\n            # Use env.step() with CallToolAction so reward flows through\n            # OpenEnv's Observation.reward — the proper RL loop contract.\n            action = CallToolAction(\n                tool_name=\"respond_to_message\",\n                arguments={\"message_id\": msg_id, \"response\": response_text},\n            )\n            step_result = env.step(action)\n            obs = step_result.observation\n\n            # Read reward from observation (populated by server's step() override)\n            reward = obs.reward if obs.reward is not None else 0.0\n            done = obs.done\n\n            # Parse tool result for error/status info\n            result_data = _extract_tool_result(obs)\n\n            if \"error\" in result_data:\n                env.call_tool(\"advance_time\", hours=1.0)\n                continue\n\n            total_reward += reward\n            target_msg = next((m for m in inbox if m[\"id\"] == msg_id), None)\n            urgency = target_msg[\"urgency\"] if target_msg else \"?\"\n            actions_taken.append({\"step\": step, \"hour\": current_hour, \"msg_id\": msg_id, \"urgency\": urgency, \"reward\": reward})\n            print(f\"  Step {step:2d} | Hour {current_hour:5.1f} | {msg_id} ({urgency:8s}) | Reward: {reward:+.1f} | Total: {total_reward:.1f}\")\n\n            if done:\n                break\n\n        final_status = json.loads(env.call_tool(\"get_status\"))\n\n    return {\"seed\": seed, \"total_reward\": total_reward, \"actions\": actions_taken, \"final_status\": final_status}\n\n\n# Run baseline on 3 seeds to get a stable estimate\nprint(\"=== PRE-TRAINING BASELINE (untrained model) ===\\n\")\nbaseline_results = []\nfor seed in [99, 42, 7]:\n    print(f\"--- Seed {seed} ---\")\n    res = evaluate_on_live_env(model, tokenizer, BASE_URL, seed=seed)\n    baseline_results.append(res)\n    print(f\"  Total: {res['total_reward']:.1f} | Actions: {len(res['actions'])}\\n\")\n\nbaseline_avg = sum(r[\"total_reward\"] for r in baseline_results) / len(baseline_results)\nprint(f\"Baseline average reward: {baseline_avg:.1f}\")\nprint(\"(Saving for comparison after training)\")",
+   "source": "# --- Pre-training baseline evaluation against live environment ---\nfrom openenv.core.env_server.mcp_types import CallToolAction\n\ndef generate_action_baseline(model, tokenizer, prompt_text):\n    \"\"\"Generate an action from the model (used for both baseline and trained eval).\"\"\"\n    msgs = [{\"role\": \"user\", \"content\": prompt_text}]\n    input_ids = tokenizer.apply_chat_template(msgs, return_tensors=\"pt\", add_generation_prompt=True)\n    if not isinstance(input_ids, torch.Tensor):\n        input_ids = input_ids[\"input_ids\"]\n    input_ids = input_ids.to(\"cuda\")\n    prompt_len = input_ids.shape[1]\n    with torch.no_grad():\n        output = model.generate(input_ids=input_ids, max_new_tokens=200, temperature=0.7,\n                                pad_token_id=tokenizer.pad_token_id, do_sample=True)\n    return tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)\n\n\ndef _extract_tool_result(obs):\n    \"\"\"Extract the JSON result dict from a CallToolObservation.\n\n    obs.result may be a plain string, a dict with 'data' key (FastMCP\n    CallToolResult serialization), or an object with a .data attribute.\n    Normalize to a dict.\n    \"\"\"\n    raw = getattr(obs, \"result\", None)\n\n    # FastMCP CallToolResult object \u2192 unwrap .data\n    if hasattr(raw, \"data\"):\n        raw = raw.data\n\n    # Serialized as {\"data\": \"...\"} dict\n    if isinstance(raw, dict) and \"data\" in raw:\n        raw = raw[\"data\"]\n\n    # Now raw should be a JSON string\n    if isinstance(raw, str):\n        try:\n            return json.loads(raw)\n        except (json.JSONDecodeError, TypeError):\n            return {}\n\n    # Already a dict\n    if isinstance(raw, dict):\n        return raw\n\n    return {}\n\n\ndef evaluate_on_live_env(model, tokenizer, base_url, seed, max_steps=20):\n    \"\"\"Run model against the live environment using OpenEnv's step() flow.\n\n    Uses env.step(CallToolAction(...)) for respond_to_message so that\n    rewards flow through Observation.reward \u2014 the standard OpenEnv contract.\n    Read-only tools (get_inbox, get_status, advance_time) use call_tool().\n    \"\"\"\n    with MCPToolClient(base_url=base_url, connect_timeout_s=60.0, message_timeout_s=120.0).sync() as env:\n        env.reset(seed=seed)\n        total_reward = 0.0\n        actions_taken = []\n\n        for step in range(max_steps):\n            inbox = json.loads(env.call_tool(\"get_inbox\"))\n            status = json.loads(env.call_tool(\"get_status\"))\n            current_hour = status[\"current_hour\"]\n\n            if status.get(\"done\"):\n                break\n\n            unhandled = [m for m in inbox if not m.get(\"handled\", False)]\n            if not unhandled:\n                env.call_tool(\"advance_time\", hours=2.0)\n                continue\n\n            completion = generate_action_baseline(model, tokenizer, build_prompt(inbox, current_hour))\n\n            match = re.search(r'respond_to_message\\s*\\(\\s*[\"\\']?(msg_\\d+)[\"\\']?\\s*,\\s*[\"\\'](.+?)[\"\\']', completion, re.DOTALL)\n            if not match:\n                id_match = re.search(r'(msg_\\d+)', completion)\n                if id_match:\n                    msg_id, response_text = id_match.group(1), completion[:200]\n                else:\n                    env.call_tool(\"advance_time\", hours=1.0)\n                    continue\n            else:\n                msg_id, response_text = match.group(1), match.group(2)\n\n            # Use env.step() with CallToolAction so reward flows through\n            # OpenEnv's Observation.reward \u2014 the proper RL loop contract.\n            action = CallToolAction(\n                tool_name=\"respond_to_message\",\n                arguments={\"message_id\": msg_id, \"response\": response_text},\n            )\n            step_result = env.step(action)\n            obs = step_result.observation\n\n            # Read reward from observation (populated by server's step() override)\n            reward = obs.reward if obs.reward is not None else 0.0\n            done = obs.done\n\n            # Parse tool result for error/status info\n            result_data = _extract_tool_result(obs)\n\n            if \"error\" in result_data:\n                env.call_tool(\"advance_time\", hours=1.0)\n                continue\n\n            total_reward += reward\n            target_msg = next((m for m in inbox if m[\"id\"] == msg_id), None)\n            urgency = target_msg[\"urgency\"] if target_msg else \"?\"\n            actions_taken.append({\"step\": step, \"hour\": current_hour, \"msg_id\": msg_id, \"urgency\": urgency, \"reward\": reward})\n            print(f\"  Step {step:2d} | Hour {current_hour:5.1f} | {msg_id} ({urgency:8s}) | Reward: {reward:+.1f} | Total: {total_reward:.1f}\")\n\n            if done:\n                break\n\n        final_status = json.loads(env.call_tool(\"get_status\"))\n\n    return {\"seed\": seed, \"total_reward\": total_reward, \"actions\": actions_taken, \"final_status\": final_status}\n\n\n# Run baseline on 3 seeds to get a stable estimate\nprint(\"=== PRE-TRAINING BASELINE (untrained model) ===\\n\")\nbaseline_results = []\nfor seed in [99, 42, 7]:\n    print(f\"--- Seed {seed} ---\")\n    res = evaluate_on_live_env(model, tokenizer, BASE_URL, seed=seed)\n    baseline_results.append(res)\n    print(f\"  Total: {res['total_reward']:.1f} | Actions: {len(res['actions'])}\\n\")\n\nbaseline_avg = sum(r[\"total_reward\"] for r in baseline_results) / len(baseline_results)\nprint(f\"Baseline average reward: {baseline_avg:.1f}\")\nprint(\"(Saving for comparison after training)\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -127,7 +241,7 @@
   {
    "cell_type": "code",
    "id": "arbp96a9wi",
-   "source": "from trl import GRPOConfig, GRPOTrainer\n\n# Build lookup from prompt text -> prompt metadata for reward scoring\n# Use first 200 chars as key (reliable — TRL may not pass custom dataset columns)\nprompt_lookup = {}\nfor p in prompts:\n    key = p[\"prompt\"][:200]\n    prompt_lookup[key] = p\n\n\ndef reward_fn(prompts, completions, **kwargs):\n    \"\"\"GRPO reward function. Scores each completion against its inbox state.\"\"\"\n    rewards = []\n    for prompt_msgs, completion in zip(prompts, completions):\n        # Extract prompt text to look up metadata\n        if isinstance(prompt_msgs, list):\n            prompt_text = prompt_msgs[-1][\"content\"] if prompt_msgs else \"\"\n        else:\n            prompt_text = str(prompt_msgs)\n\n        key = prompt_text[:200]\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            if completion and isinstance(completion[0], (int, float)):\n                comp_text = tokenizer.decode(completion, skip_special_tokens=True)\n            else:\n                comp_text = completion[-1].get(\"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\ntraining_args = GRPOConfig(\n    output_dir=\"crisisinbox-grpo-output\",\n    num_train_epochs=3,\n    per_device_train_batch_size=2,\n    gradient_accumulation_steps=2,\n    learning_rate=1e-5,\n    max_completion_length=256,\n    max_prompt_length=1024,\n    num_generations=4,\n    logging_steps=1,\n    save_steps=100,\n    bf16=_use_bf16,\n    fp16=not _use_bf16,\n    sync_ref_model=True,\n)\n\n# Let GRPOTrainer handle PEFT wrapping (avoids dtype mismatches from manual setup)\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_fn,\n    args=training_args,\n    train_dataset=dataset,\n    peft_config=lora_config,\n)\n\n# After trainer init, update model ref to the PEFT-wrapped version\nmodel = trainer.model\n\nprint(f\"Trainer configured — {len(prompt_lookup)} unique prompt keys\")\nprint(f\"Precision: {'bf16' if _use_bf16 else 'fp16'}\")\nprint(f\"Training for {training_args.num_train_epochs} epochs\")\nprint(\"Ready to train\")",
+   "source": "from trl import GRPOConfig, GRPOTrainer\n\n# Build lookup from prompt text -> prompt metadata for reward scoring\n# Use first 200 chars as key (reliable \u2014 TRL may not pass custom dataset columns)\nprompt_lookup = {}\nfor p in prompts:\n    key = p[\"prompt\"][:200]\n    prompt_lookup[key] = p\n\n\ndef reward_fn(prompts, completions, **kwargs):\n    \"\"\"GRPO reward function. Scores each completion against its inbox state.\"\"\"\n    rewards = []\n    for prompt_msgs, completion in zip(prompts, completions):\n        # Extract prompt text to look up metadata\n        if isinstance(prompt_msgs, list):\n            prompt_text = prompt_msgs[-1][\"content\"] if prompt_msgs else \"\"\n        else:\n            prompt_text = str(prompt_msgs)\n\n        key = prompt_text[:200]\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            if completion and isinstance(completion[0], (int, float)):\n                comp_text = tokenizer.decode(completion, skip_special_tokens=True)\n            else:\n                comp_text = completion[-1].get(\"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\ntraining_args = GRPOConfig(\n    output_dir=\"crisisinbox-grpo-output\",\n    num_train_epochs=3,\n    per_device_train_batch_size=2,\n    gradient_accumulation_steps=2,\n    learning_rate=1e-5,\n    max_completion_length=256,\n    max_prompt_length=1024,\n    num_generations=4,\n    logging_steps=1,\n    save_steps=100,\n    bf16=_use_bf16,\n    fp16=not _use_bf16,\n    sync_ref_model=True,\n)\n\n# Let GRPOTrainer handle PEFT wrapping (avoids dtype mismatches from manual setup)\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=reward_fn,\n    args=training_args,\n    train_dataset=dataset,\n    peft_config=lora_config,\n)\n\n# After trainer init, update model ref to the PEFT-wrapped version\nmodel = trainer.model\n\nprint(f\"Trainer configured \u2014 {len(prompt_lookup)} unique prompt keys\")\nprint(f\"Precision: {'bf16' if _use_bf16 else 'fp16'}\")\nprint(f\"Training for {training_args.num_train_epochs} epochs\")\nprint(\"Ready to train\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []

rewards.py

@@ -0,0 +1,107 @@
+"""
+CrisisInbox Reward Function — Single Source of Truth.
+
+Both the environment server and training notebooks import from here.
+No OpenEnv dependencies — only requires models.py (pydantic).
+"""
+
+try:
+    from .models import Message, Urgency
+except ImportError:
+    from models import Message, Urgency
+
+
+_FAMILY_SENDERS = {"mom", "sister", "neighbor dave", "emma"}
+_FAMILY_TONE_WORDS = {"love", "safe", "worried", "sorry", "care", "okay", "miss", "hang in there"}
+_FORMAL_TONE_WORDS = {"confirm", "attached", "documentation", "regarding", "request", "please", "submit"}
+
+
+def tone_multiplier(sender: str, response: str) -> float:
+    """
+    Small reward multiplier for tone-appropriate responses.
+
+    Family/personal senders reward empathetic language.
+    Professional/institutional senders reward formal language.
+    Returns 1.0-1.15 (bonus) or 1.0 (neutral). Never penalizes.
+    """
+    resp_lower = response.lower()
+    sender_lower = sender.lower()
+
+    if any(f in sender_lower for f in _FAMILY_SENDERS):
+        matches = sum(1 for w in _FAMILY_TONE_WORDS if w in resp_lower)
+        if matches >= 2:
+            return 1.15
+        elif matches == 1:
+            return 1.07
+    else:
+        matches = sum(1 for w in _FORMAL_TONE_WORDS if w in resp_lower)
+        if matches >= 2:
+            return 1.1
+        elif matches == 1:
+            return 1.05
+
+    return 1.0
+
+
+def calculate_reward(
+    msg: Message,
+    current_hour: float,
+    response: str,
+    superseded: dict[str, str],
+    visible_messages: list[Message] | None = None,
+    handled: dict[str, str] | None = None,
+) -> float:
+    """
+    Calculate reward for handling a message.
+
+    Reward signals:
+    - Base reward by urgency (critical=10, high=5, medium=3, low=1)
+    - Deadline timing bonus (up to +50% for early, -75% for late)
+    - Response quality (penalty for very short responses)
+    - Tone awareness (up to +15% for matching tone to sender type)
+    - Schema drift adaptation bonus (+50% for handling drift messages)
+    - Penalty for acting on superseded/stale information (-50%)
+    - Priority penalty (-70% for choosing low/medium when critical is pending)
+    """
+    base_rewards = {
+        Urgency.CRITICAL: 10.0,
+        Urgency.HIGH: 5.0,
+        Urgency.MEDIUM: 3.0,
+        Urgency.LOW: 1.0,
+    }
+    reward = base_rewards.get(msg.urgency, 1.0)
+
+    # Deadline timing
+    if msg.deadline_hours is not None:
+        if current_hour <= msg.deadline_hours:
+            time_remaining_frac = (msg.deadline_hours - current_hour) / max(msg.deadline_hours, 1.0)
+            reward *= 1.0 + 0.5 * time_remaining_frac
+        else:
+            reward *= 0.25
+
+    # Response quality - penalty for very short/empty responses
+    if len(response.strip()) < 10:
+        reward *= 0.5
+
+    # Tone awareness: small bonus for matching tone to sender type
+    reward *= tone_multiplier(msg.sender, response)
+
+    # Drift adaptation bonus
+    if msg.drift_flag:
+        reward *= 1.5
+
+    # Penalty for responding to a superseded message (stale info)
+    if msg.id in superseded:
+        reward *= 0.5
+
+    # Priority penalty: choosing low/medium when unhandled critical messages exist
+    if visible_messages and handled is not None:
+        has_unhandled_critical = any(
+            m.urgency == Urgency.CRITICAL and m.id not in handled
+            for m in visible_messages
+            if m.id != msg.id  # exclude the message being handled now
+        )
+        if has_unhandled_critical and msg.urgency in (Urgency.LOW, Urgency.MEDIUM):
+            reward *= 0.3
+
+    return round(reward, 2)

server/crisis_inbox_environment.py

@@ -25,10 +25,12 @@ try:
     from ..models import Channel, Message, Urgency
     from ..messages import ALL_MESSAGES
     from ..drift_events import ALL_DRIFT_EVENTS, DriftEvent, select_drift_events
+    from ..rewards import calculate_reward
 except ImportError:
     from models import Channel, Message, Urgency
     from messages import ALL_MESSAGES
     from drift_events import ALL_DRIFT_EVENTS, DriftEvent, select_drift_events
+    from rewards import calculate_reward
 
 
 class CrisisInboxEnvironment(MCPEnvironment):
@@ -161,7 +163,7 @@ class CrisisInboxEnvironment(MCPEnvironment):
                 })
 
             # Calculate reward
-            reward = _calculate_reward(
+            reward = calculate_reward(
                 msg, self._current_hour, response, self._superseded,
                 self._visible_messages, self._handled,
             )
@@ -397,63 +399,3 @@ class CrisisInboxEnvironment(MCPEnvironment):
     @property
     def state(self) -> State:
         return self._state
-
-
-def _calculate_reward(
-    msg: Message,
-    current_hour: float,
-    response: str,
-    superseded: dict[str, str],
-    visible_messages: list[Message] | None = None,
-    handled: dict[str, str] | None = None,
-) -> float:
-    """
-    Calculate reward for handling a message.
-
-    Reward signals:
-    - Base reward by urgency (critical=10, high=5, medium=3, low=1)
-    - Deadline timing bonus (up to +50% for early, -75% for late)
-    - Response quality (penalty for very short responses)
-    - Schema drift adaptation bonus (+50% for handling drift messages)
-    - Penalty for acting on superseded/stale information (-50%)
-    - Priority penalty (-70% for choosing low/medium when critical is pending)
-    """
-    base_rewards = {
-        Urgency.CRITICAL: 10.0,
-        Urgency.HIGH: 5.0,
-        Urgency.MEDIUM: 3.0,
-        Urgency.LOW: 1.0,
-    }
-    reward = base_rewards.get(msg.urgency, 1.0)
-
-    # Deadline timing
-    if msg.deadline_hours is not None:
-        if current_hour <= msg.deadline_hours:
-            time_remaining_frac = (msg.deadline_hours - current_hour) / max(msg.deadline_hours, 1.0)
-            reward *= 1.0 + 0.5 * time_remaining_frac
-        else:
-            reward *= 0.25
-
-    # Response quality - penalty for very short/empty responses
-    if len(response.strip()) < 10:
-        reward *= 0.5
-
-    # Drift adaptation bonus
-    if msg.drift_flag:
-        reward *= 1.5
-
-    # Penalty for responding to a superseded message (stale info)
-    if msg.id in superseded:
-        reward *= 0.5
-
-    # Priority penalty: choosing low/medium when unhandled critical messages exist
-    if visible_messages and handled is not None:
-        has_unhandled_critical = any(
-            m.urgency == Urgency.CRITICAL and m.id not in handled
-            for m in visible_messages
-            if m.id != msg.id  # exclude the message being handled now
-        )
-        if has_unhandled_critical and msg.urgency in (Urgency.LOW, Urgency.MEDIUM):
-            reward *= 0.3
-
-    return round(reward, 2)