Upload train_flatmate_rl_trl.ipynb with huggingface_hub

train_flatmate_rl_trl.ipynb

@@ -0,0 +1,777 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "3905a08b",
+   "metadata": {},
+   "source": [
+    "# Train a Flatmate RL Action Policy with TRL\n",
+    "\n",
+    "This notebook connects to the Hugging Face Space endpoint, collects rollout examples over OpenEnv websocket sessions, and fine-tunes a small causal language model to emit Flatmate RL JSON actions. The training path uses TRL `SFTTrainer`, which is the most stable starting point for this mixed natural-language plus structured-tool action space.\n",
+    "\n",
+    "Endpoint used here: `https://huggingface.co/spaces/kushalExplores/flatmate_rl`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54f0ddc0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Install notebook dependencies. Restart the kernel after this cell if Colab/Jupyter asks you to.\n",
+    "%pip install -q \"trl>=0.23.0\" \"transformers>=4.46.0\" accelerate datasets peft websockets huggingface_hub matplotlib pandas"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6a37c34",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "import asyncio\n",
+    "import json\n",
+    "import random\n",
+    "from dataclasses import dataclass\n",
+    "from pathlib import Path\n",
+    "from typing import Any\n",
+    "from urllib.parse import urlparse\n",
+    "\n",
+    "import websockets\n",
+    "from datasets import Dataset\n",
+    "\n",
+    "SPACE_HTTP_URL = \"https://kushalexplores-flatmate-rl.hf.space\"\n",
+    "SCENARIOS = [\n",
+    "    \"task_visit_single\",\n",
+    "    \"task_visit_single_hidden_flex\",\n",
+    "    \"task_visit_multi\",\n",
+    "    \"task_visit_single_seller_followup\",\n",
+    "]\n",
+    "\n",
+    "def ws_url_from_http(base_url: str) -> str:\n",
+    "    parsed = urlparse(base_url.rstrip(\"/\"))\n",
+    "    scheme = \"wss\" if parsed.scheme == \"https\" else \"ws\"\n",
+    "    return f\"{scheme}://{parsed.netloc}/ws\"\n",
+    "\n",
+    "SPACE_WS_URL = ws_url_from_http(SPACE_HTTP_URL)\n",
+    "SPACE_WS_URL"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3e10f23e",
+   "metadata": {},
+   "source": [
+    "## Endpoint Client\n",
+    "\n",
+    "OpenEnv's plain HTTP `/reset` and `/step` endpoints are stateless. Use `/ws` for multi-step episodes because the websocket session keeps one environment instance alive across reset and step calls."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f958cca7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class FlatmateEndpoint:\n",
+    "    def __init__(self, ws_url: str = SPACE_WS_URL, timeout_s: float = 120.0):\n",
+    "        self.ws_url = ws_url\n",
+    "        self.timeout_s = timeout_s\n",
+    "\n",
+    "    async def __aenter__(self):\n",
+    "        self.ws = await websockets.connect(self.ws_url, open_timeout=self.timeout_s, ping_timeout=self.timeout_s)\n",
+    "        return self\n",
+    "\n",
+    "    async def __aexit__(self, exc_type, exc, tb):\n",
+    "        try:\n",
+    "            await self.ws.send(json.dumps({\"type\": \"close\"}))\n",
+    "        finally:\n",
+    "            await self.ws.close()\n",
+    "\n",
+    "    async def _send(self, payload: dict[str, Any]) -> dict[str, Any]:\n",
+    "        await self.ws.send(json.dumps(payload))\n",
+    "        raw = await asyncio.wait_for(self.ws.recv(), timeout=self.timeout_s)\n",
+    "        message = json.loads(raw)\n",
+    "        if message.get(\"type\") == \"error\":\n",
+    "            raise RuntimeError(message.get(\"data\", message))\n",
+    "        data = message[\"data\"]\n",
+    "        obs = data.get(\"observation\", {})\n",
+    "        obs[\"reward\"] = data.get(\"reward\")\n",
+    "        obs[\"done\"] = data.get(\"done\", False)\n",
+    "        return obs\n",
+    "\n",
+    "    async def reset(self, scenario_id: str, seed: int | None = None) -> dict[str, Any]:\n",
+    "        data: dict[str, Any] = {\"scenario_id\": scenario_id}\n",
+    "        if seed is not None:\n",
+    "            data[\"seed\"] = seed\n",
+    "        return await self._send({\"type\": \"reset\", \"data\": data})\n",
+    "\n",
+    "    async def step(self, action: dict[str, Any]) -> dict[str, Any]:\n",
+    "        return await self._send({\"type\": \"step\", \"data\": action})\n",
+    "\n",
+    "async def smoke_test_endpoint():\n",
+    "    async with FlatmateEndpoint() as env:\n",
+    "        obs = await env.reset(\"task_visit_single\", seed=1)\n",
+    "        print(obs[\"scenario_id\"], obs[\"status\"])\n",
+    "        print(obs.get(\"last_user_message\") or obs.get(\"current_user_request\"))\n",
+    "\n",
+    "await smoke_test_endpoint()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fe2ad079",
+   "metadata": {},
+   "source": [
+    "## Rollout Policy for Data Collection\n",
+    "\n",
+    "This heuristic is intentionally simple. It produces valid-looking action examples from endpoint observations; after SFT, replace it with model generation and keep the same evaluator."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "611b1ac4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def tool_names(obs: dict[str, Any]) -> list[str]:\n",
+    "    return [str(t.get(\"tool\", t.get(\"tool_name\", \"\"))) for t in obs.get(\"tool_trace\", [])]\n",
+    "\n",
+    "def action_policy(obs: dict[str, Any]) -> dict[str, Any] | None:\n",
+    "    tools = tool_names(obs)\n",
+    "    phase = obs.get(\"phase\", \"buyer\")\n",
+    "    remaining = set(obs.get(\"remaining_required_fields\", []))\n",
+    "    scenario_id = obs.get(\"scenario_id\", \"task_visit_single\")\n",
+    "\n",
+    "    if phase == \"seller\" and not obs.get(\"seller_profile_stored\"):\n",
+    "        if remaining:\n",
+    "            return {\"action_type\": \"assistant_message\", \"assistant_message\": \"Please share the household dietary setup, who the flat is for, and available visit time slots.\"}\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"store_seller_details\", \"tool_arguments\": {}}\n",
+    "\n",
+    "    if not obs.get(\"buyer_profile_stored\"):\n",
+    "        if \"diet\" in remaining and \"visit_availability\" in remaining:\n",
+    "            return {\"action_type\": \"assistant_message\", \"assistant_message\": \"Please share your dietary preference and visit availability.\"}\n",
+    "        if \"diet\" in remaining:\n",
+    "            return {\"action_type\": \"assistant_message\", \"assistant_message\": \"Please share your dietary preference.\"}\n",
+    "        if \"visit_availability\" in remaining:\n",
+    "            return {\"action_type\": \"assistant_message\", \"assistant_message\": \"Please share your visit availability.\"}\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"store_user_details\", \"tool_arguments\": {}}\n",
+    "\n",
+    "    if \"search_posts\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"search_posts\", \"tool_arguments\": {}}\n",
+    "\n",
+    "    post_ids = [\"post_031\", \"post_052\"] if scenario_id == \"task_visit_multi\" else [\"post_031\"]\n",
+    "    if \"match_location_preference\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"match_location_preference\", \"tool_arguments\": {\"post_ids\": post_ids}}\n",
+    "    if \"get_commute_time\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"get_commute_time\", \"tool_arguments\": {\"post_ids\": post_ids}}\n",
+    "    if \"check_calendar_slots\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"check_calendar_slots\", \"tool_arguments\": {\"post_ids\": post_ids}}\n",
+    "    if \"shortlist\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"shortlist\", \"tool_arguments\": {\"post_ids\": post_ids}}\n",
+    "    if \"contact_poster\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"contact_poster\", \"tool_arguments\": {\"post_id\": post_ids[0], \"time_text\": \"tomorrow 7pm\"}}\n",
+    "    if \"book_viewing\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"book_viewing\", \"tool_arguments\": {\"post_id\": post_ids[0], \"time_text\": \"tomorrow 7pm\"}}\n",
+    "\n",
+    "    return None\n",
+    "\n",
+    "def flatten_observation(obs: dict[str, Any]) -> str:\n",
+    "    visible = {\n",
+    "        \"scenario_id\": obs.get(\"scenario_id\"),\n",
+    "        \"phase\": obs.get(\"phase\"),\n",
+    "        \"status\": obs.get(\"status\"),\n",
+    "        \"last_user_message\": obs.get(\"last_user_message\"),\n",
+    "        \"current_user_request\": obs.get(\"current_user_request\"),\n",
+    "        \"available_tools\": obs.get(\"available_tools\", []),\n",
+    "        \"remaining_required_fields\": obs.get(\"remaining_required_fields\", []),\n",
+    "        \"prerequisites_satisfied\": obs.get(\"prerequisites_satisfied\", {}),\n",
+    "        \"recent_tool_calls\": obs.get(\"recent_tool_calls\", []),\n",
+    "        \"last_tool_result\": obs.get(\"last_tool_result\", {}),\n",
+    "        \"violations\": obs.get(\"violations\", []),\n",
+    "        \"booked_visits\": obs.get(\"booked_visits\", []),\n",
+    "        \"feedback_summary\": obs.get(\"feedback_summary\", \"\"),\n",
+    "    }\n",
+    "    return json.dumps(visible, ensure_ascii=False, sort_keys=True)\n",
+    "\n",
+    "def make_training_text(obs: dict[str, Any], action: dict[str, Any]) -> str:\n",
+    "    return (\n",
+    "        \"You are a broker policy for the Flatmate RL environment. \"\n",
+    "        \"Given an observation, return exactly one JSON action.\\n\\n\"\n",
+    "        f\"Observation:\\n{flatten_observation(obs)}\\n\\n\"\n",
+    "        f\"Action:\\n{json.dumps(action, ensure_ascii=False, sort_keys=True)}\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7b22fa13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@dataclass\n",
+    "class RolloutConfig:\n",
+    "    train_episodes_per_task: int = 4\n",
+    "    test_episodes_per_task: int = 2\n",
+    "    max_steps: int = 20\n",
+    "    seed: int = 7\n",
+    "\n",
+    "async def collect_one_episode(\n",
+    "    scenario_id: str,\n",
+    "    episode_id: str,\n",
+    "    episode_idx: int,\n",
+    "    split: str,\n",
+    "    seed: int,\n",
+    "    max_steps: int,\n",
+    ") -> list[dict[str, Any]]:\n",
+    "    rows: list[dict[str, Any]] = []\n",
+    "    async with FlatmateEndpoint() as env:\n",
+    "        obs = await env.reset(scenario_id, seed=seed)\n",
+    "        total_reward = 0.0\n",
+    "        for step_idx in range(max_steps):\n",
+    "            action = action_policy(obs)\n",
+    "            if action is None or obs.get(\"done\"):\n",
+    "                break\n",
+    "            rows.append({\n",
+    "                \"text\": make_training_text(obs, action),\n",
+    "                \"episode_id\": episode_id,\n",
+    "                \"episode_idx\": episode_idx,\n",
+    "                \"split\": split,\n",
+    "                \"scenario_id\": scenario_id,\n",
+    "                \"seed\": seed,\n",
+    "                \"step\": step_idx,\n",
+    "                \"action\": json.dumps(action, sort_keys=True),\n",
+    "            })\n",
+    "            obs = await env.step(action)\n",
+    "            total_reward += float(obs.get(\"reward\") or obs.get(\"step_reward\") or 0.0)\n",
+    "            if obs.get(\"done\"):\n",
+    "                break\n",
+    "    print(f\"split={split:5s} episode={episode_id} scenario={scenario_id} rows={len(rows)} total_reward={total_reward:.2f}\")\n",
+    "    return rows\n",
+    "\n",
+    "async def collect_balanced_rollouts(config: RolloutConfig = RolloutConfig()) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:\n",
+    "    train_rows: list[dict[str, Any]] = []\n",
+    "    test_rows: list[dict[str, Any]] = []\n",
+    "    episode_idx = 0\n",
+    "\n",
+    "    for scenario_idx, scenario_id in enumerate(SCENARIOS):\n",
+    "        for task_episode_idx in range(config.train_episodes_per_task):\n",
+    "            seed = config.seed + scenario_idx * 100 + task_episode_idx\n",
+    "            episode_id = f\"train_{scenario_id}_{task_episode_idx:03d}\"\n",
+    "            train_rows.extend(await collect_one_episode(scenario_id, episode_id, episode_idx, \"train\", seed, config.max_steps))\n",
+    "            episode_idx += 1\n",
+    "\n",
+    "        for task_episode_idx in range(config.test_episodes_per_task):\n",
+    "            seed = 900 + config.seed + scenario_idx * 100 + task_episode_idx\n",
+    "            episode_id = f\"test_{scenario_id}_{task_episode_idx:03d}\"\n",
+    "            test_rows.extend(await collect_one_episode(scenario_id, episode_id, episode_idx, \"test\", seed, config.max_steps))\n",
+    "            episode_idx += 1\n",
+    "\n",
+    "    return train_rows, test_rows\n",
+    "\n",
+    "print(\"Note: seeded resets create value variants while preserving the same episode flow. Upload the updated Space before using this against the hosted endpoint.\")\n",
+    "train_rows, test_rows = await collect_balanced_rollouts(\n",
+    "    RolloutConfig(train_episodes_per_task=4, test_episodes_per_task=2, max_steps=20, seed=7)\n",
+    ")\n",
+    "rows = train_rows + test_rows\n",
+    "dataset = Dataset.from_list(rows)\n",
+    "train_dataset = Dataset.from_list(train_rows)\n",
+    "test_dataset = Dataset.from_list(test_rows)\n",
+    "\n",
+    "print({\n",
+    "    \"train_rows\": len(train_dataset),\n",
+    "    \"test_rows\": len(test_dataset),\n",
+    "    \"total_rows\": len(dataset),\n",
+    "    \"train_episodes\": len(set(train_dataset[\"episode_id\"])),\n",
+    "    \"test_episodes\": len(set(test_dataset[\"episode_id\"])),\n",
+    "})\n",
+    "print(\"train scenarios\", sorted(set(train_dataset[\"scenario_id\"])))\n",
+    "print(\"test scenarios\", sorted(set(test_dataset[\"scenario_id\"])))\n",
+    "print(\"train episodes by scenario\")\n",
+    "display(pd.DataFrame(train_rows).groupby(\"scenario_id\")[\"episode_id\"].nunique().rename(\"episodes\"))\n",
+    "print(\"test episodes by scenario\")\n",
+    "display(pd.DataFrame(test_rows).groupby(\"scenario_id\")[\"episode_id\"].nunique().rename(\"episodes\"))\n",
+    "{\"train\": train_dataset, \"test\": test_dataset}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "665b46fa",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import LoraConfig\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
+    "from trl import SFTConfig, SFTTrainer\n",
+    "\n",
+    "MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
+    "OUTPUT_DIR = \"flatmate-rl-action-policy\"\n",
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)\n",
+    "if tokenizer.pad_token is None:\n",
+    "    tokenizer.pad_token = tokenizer.eos_token\n",
+    "\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    MODEL_NAME,\n",
+    "    trust_remote_code=True,\n",
+    "    device_map=\"auto\",\n",
+    ")\n",
+    "model.config.use_cache = False\n",
+    "\n",
+    "peft_config = LoraConfig(\n",
+    "    r=16,\n",
+    "    lora_alpha=32,\n",
+    "    lora_dropout=0.05,\n",
+    "    bias=\"none\",\n",
+    "    task_type=\"CAUSAL_LM\",\n",
+    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+    ")\n",
+    "\n",
+    "training_args = SFTConfig(\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    dataset_text_field=\"text\",\n",
+    "    max_length=1536,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=8,\n",
+    "    num_train_epochs=1,\n",
+    "    learning_rate=5e-5,\n",
+    "    logging_steps=5,\n",
+    "    save_steps=50,\n",
+    "    save_total_limit=2,\n",
+    "    packing=False,\n",
+    "    report_to=\"none\",\n",
+    ")\n",
+    "\n",
+    "trainer = SFTTrainer(\n",
+    "    model=model,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_dataset,\n",
+    "    eval_dataset=test_dataset,\n",
+    "    processing_class=tokenizer,\n",
+    "    peft_config=peft_config,\n",
+    ")\n",
+    "\n",
+    "train_result = trainer.train()\n",
+    "test_metrics = trainer.evaluate(eval_dataset=test_dataset)\n",
+    "train_log_history = trainer.state.log_history\n",
+    "trainer.save_model(OUTPUT_DIR)\n",
+    "tokenizer.save_pretrained(OUTPUT_DIR)\n",
+    "print(\"heldout_test_metrics\", test_metrics)\n",
+    "train_result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "22d9fc14",
+   "metadata": {},
+   "source": [
+    "## Training Log\n",
+    "\n",
+    "Plot the logged training loss over optimizer steps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c3e44d74",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import pandas as pd\n",
+    "\n",
+    "log_path = Path(OUTPUT_DIR) / \"train_log_history.json\"\n",
+    "log_path.parent.mkdir(parents=True, exist_ok=True)\n",
+    "log_path.write_text(json.dumps(train_log_history, indent=2))\n",
+    "\n",
+    "def plot_training_log(log_history, title: str = \"SFT training loss\"):\n",
+    "    rows = [row for row in log_history if \"loss\" in row and \"step\" in row]\n",
+    "    if not rows:\n",
+    "        print(\"No loss rows found in trainer.state.log_history yet.\")\n",
+    "        return None\n",
+    "    df = pd.DataFrame(rows)\n",
+    "    ax = df.plot(x=\"step\", y=\"loss\", marker=\"o\", figsize=(7, 4), title=title)\n",
+    "    ax.set_xlabel(\"optimizer step\")\n",
+    "    ax.set_ylabel(\"loss\")\n",
+    "    ax.grid(True, alpha=0.3)\n",
+    "    plt.show()\n",
+    "    return df\n",
+    "\n",
+    "train_log_df = plot_training_log(train_log_history)\n",
+    "train_log_df.tail() if train_log_df is not None else None"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "539548f7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from peft import AutoPeftModelForCausalLM\n",
+    "\n",
+    "# Load both the base model and the saved fine-tuned adapter from disk for comparison.\n",
+    "try:\n",
+    "    del model\n",
+    "except NameError:\n",
+    "    pass\n",
+    "\n",
+    "base_model_for_eval = AutoModelForCausalLM.from_pretrained(\n",
+    "    MODEL_NAME,\n",
+    "    trust_remote_code=True,\n",
+    "    device_map=\"auto\",\n",
+    ")\n",
+    "base_model_for_eval.eval()\n",
+    "base_model_for_eval.config.use_cache = False\n",
+    "\n",
+    "loaded_model_for_eval = AutoPeftModelForCausalLM.from_pretrained(OUTPUT_DIR, device_map=\"auto\")\n",
+    "loaded_model_for_eval.eval()\n",
+    "loaded_model_for_eval.config.use_cache = False\n",
+    "active_model = loaded_model_for_eval\n",
+    "print(f\"Loaded base model from {MODEL_NAME}\")\n",
+    "print(f\"Loaded saved SFT model from {OUTPUT_DIR}\")\n",
+    "\n",
+    "TEST_SEEDS = (901, 902)\n",
+    "\n",
+    "\n",
+    "def prompt_from_observation(obs: dict[str, Any]) -> str:\n",
+    "    return (\n",
+    "        \"You are a broker policy for the Flatmate RL environment. \"\n",
+    "        \"Given an observation, return exactly one JSON action.\\n\\n\"\n",
+    "        f\"Observation:\\n{flatten_observation(obs)}\\n\\nAction:\\n\"\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "def _first_balanced_json(text: str) -> str:\n",
+    "    start = text.find(\"{\")\n",
+    "    if start == -1:\n",
+    "        raise ValueError(f\"No JSON object found in generation: {text!r}\")\n",
+    "    depth = 0\n",
+    "    in_string = False\n",
+    "    escape = False\n",
+    "    for index, char in enumerate(text[start:], start=start):\n",
+    "        if escape:\n",
+    "            escape = False\n",
+    "            continue\n",
+    "        if char == \"\\\\\" and in_string:\n",
+    "            escape = True\n",
+    "            continue\n",
+    "        if char == '\\\"':\n",
+    "            in_string = not in_string\n",
+    "            continue\n",
+    "        if in_string:\n",
+    "            continue\n",
+    "        if char == \"{\":\n",
+    "            depth += 1\n",
+    "        elif char == \"}\":\n",
+    "            depth -= 1\n",
+    "            if depth == 0:\n",
+    "                return text[start : index + 1]\n",
+    "    raise ValueError(f\"Unterminated JSON object in generation: {text!r}\")\n",
+    "\n",
+    "\n",
+    "def normalize_action(action: dict[str, Any]) -> dict[str, Any]:\n",
+    "    if action.get(\"action_type\") == \"assistant_message\" and str(action.get(\"assistant_message\", \"\")).strip():\n",
+    "        return {\n",
+    "            \"action_type\": \"assistant_message\",\n",
+    "            \"assistant_message\": str(action[\"assistant_message\"]),\n",
+    "        }\n",
+    "    if action.get(\"action_type\") == \"tool_call\" and str(action.get(\"tool_name\", \"\")).strip():\n",
+    "        tool_arguments = action.get(\"tool_arguments\", {})\n",
+    "        return {\n",
+    "            \"action_type\": \"tool_call\",\n",
+    "            \"tool_name\": str(action[\"tool_name\"]),\n",
+    "            \"tool_arguments\": tool_arguments if isinstance(tool_arguments, dict) else {},\n",
+    "        }\n",
+    "    raise ValueError(f\"Invalid action shape: {action!r}\")\n",
+    "\n",
+    "\n",
+    "def parse_action(text: str) -> dict[str, Any]:\n",
+    "    return normalize_action(json.loads(_first_balanced_json(text)))\n",
+    "\n",
+    "\n",
+    "def heuristic_policy(obs: dict[str, Any]) -> dict[str, Any]:\n",
+    "    action = action_policy(obs)\n",
+    "    if action is None:\n",
+    "        return {\"action_type\": \"assistant_message\", \"assistant_message\": \"Could you confirm the details needed for scheduling?\"}\n",
+    "    return action\n",
+    "\n",
+    "\n",
+    "def raw_generate_action_text(obs: dict[str, Any]) -> str:\n",
+    "    prompt = prompt_from_observation(obs) + \"{\"\n",
+    "    inputs = tokenizer(prompt, return_tensors=\"pt\").to(active_model.device)\n",
+    "    active_model.generation_config.do_sample = False\n",
+    "    active_model.generation_config.temperature = None\n",
+    "    active_model.generation_config.top_p = None\n",
+    "    active_model.generation_config.top_k = None\n",
+    "    with torch.no_grad():\n",
+    "        output = active_model.generate(\n",
+    "            **inputs,\n",
+    "            max_new_tokens=80,\n",
+    "            do_sample=False,\n",
+    "            repetition_penalty=1.15,\n",
+    "            no_repeat_ngram_size=3,\n",
+    "            eos_token_id=tokenizer.eos_token_id,\n",
+    "            pad_token_id=tokenizer.eos_token_id,\n",
+    "        )\n",
+    "    return \"{\" + tokenizer.decode(output[0][inputs[\"input_ids\"].shape[-1]:], skip_special_tokens=True)\n",
+    "\n",
+    "\n",
+    "def model_action_or_error(obs: dict[str, Any]) -> tuple[dict[str, Any] | None, str, str]:\n",
+    "    raw = raw_generate_action_text(obs)\n",
+    "    try:\n",
+    "        return parse_action(raw), raw, \"\"\n",
+    "    except Exception as exc:\n",
+    "        return None, raw, str(exc)\n",
+    "\n",
+    "\n",
+    "async def sanity_check_generations(model_label: str, limit: int = 4):\n",
+    "    rows = []\n",
+    "    for scenario_id in SCENARIOS[:limit]:\n",
+    "        async with FlatmateEndpoint() as env:\n",
+    "            obs = await env.reset(scenario_id, seed=TEST_SEEDS[0])\n",
+    "            action, raw, error = model_action_or_error(obs)\n",
+    "            rows.append({\n",
+    "                \"model\": model_label,\n",
+    "                \"scenario_id\": scenario_id,\n",
+    "                \"json_ok\": action is not None,\n",
+    "                \"raw\": raw[:240],\n",
+    "                \"parsed_action\": action,\n",
+    "                \"error\": error,\n",
+    "            })\n",
+    "    return pd.DataFrame(rows)\n",
+    "\n",
+    "\n",
+    "async def evaluate_heuristic(label: str = \"heuristic\", scenarios=SCENARIOS, seeds=TEST_SEEDS, max_steps: int = 20, verbose: bool = False):\n",
+    "    rows = []\n",
+    "    for scenario_id in scenarios:\n",
+    "        for seed in seeds:\n",
+    "            async with FlatmateEndpoint() as env:\n",
+    "                obs = await env.reset(scenario_id, seed=seed)\n",
+    "                total_reward = 0.0\n",
+    "                steps = 0\n",
+    "                for step_idx in range(max_steps):\n",
+    "                    action = heuristic_policy(obs)\n",
+    "                    if verbose:\n",
+    "                        print(label, scenario_id, seed, step_idx, action)\n",
+    "                    obs = await env.step(action)\n",
+    "                    steps = step_idx + 1\n",
+    "                    total_reward += float(obs.get(\"reward\") or obs.get(\"step_reward\") or 0.0)\n",
+    "                    if obs.get(\"done\"):\n",
+    "                        break\n",
+    "                rows.append({\n",
+    "                    \"policy\": label,\n",
+    "                    \"scenario_id\": scenario_id,\n",
+    "                    \"seed\": seed,\n",
+    "                    \"total_reward\": total_reward,\n",
+    "                    \"done\": bool(obs.get(\"done\")),\n",
+    "                    \"bookings\": len(obs.get(\"booked_visits\", [])),\n",
+    "                    \"violations\": len(obs.get(\"violations\", [])),\n",
+    "                    \"steps\": steps,\n",
+    "                    \"parse_errors\": 0,\n",
+    "                })\n",
+    "    return rows\n",
+    "\n",
+    "\n",
+    "async def evaluate_model_policy(label: str, scenarios=SCENARIOS, seeds=TEST_SEEDS, max_steps: int = 20, verbose: bool = False):\n",
+    "    rows = []\n",
+    "    for scenario_id in scenarios:\n",
+    "        for seed in seeds:\n",
+    "            async with FlatmateEndpoint() as env:\n",
+    "                obs = await env.reset(scenario_id, seed=seed)\n",
+    "                total_reward = 0.0\n",
+    "                steps = 0\n",
+    "                parse_errors = 0\n",
+    "                last_error = \"\"\n",
+    "                for step_idx in range(max_steps):\n",
+    "                    action, raw, error = model_action_or_error(obs)\n",
+    "                    if action is None:\n",
+    "                        parse_errors += 1\n",
+    "                        last_error = error\n",
+    "                        if verbose:\n",
+    "                            print(label, scenario_id, seed, f\"step={step_idx:02d}\", \"PARSE_ERROR\", raw[:220])\n",
+    "                        total_reward -= 1.0\n",
+    "                        break\n",
+    "                    if verbose:\n",
+    "                        print(label, scenario_id, seed, f\"step={step_idx:02d}\", action)\n",
+    "                    obs = await env.step(action)\n",
+    "                    steps = step_idx + 1\n",
+    "                    total_reward += float(obs.get(\"reward\") or obs.get(\"step_reward\") or 0.0)\n",
+    "                    if obs.get(\"done\"):\n",
+    "                        break\n",
+    "                rows.append({\n",
+    "                    \"policy\": label,\n",
+    "                    \"scenario_id\": scenario_id,\n",
+    "                    \"seed\": seed,\n",
+    "                    \"total_reward\": total_reward,\n",
+    "                    \"done\": bool(obs.get(\"done\")),\n",
+    "                    \"bookings\": len(obs.get(\"booked_visits\", [])),\n",
+    "                    \"violations\": len(obs.get(\"violations\", [])),\n",
+    "                    \"steps\": steps,\n",
+    "                    \"parse_errors\": parse_errors,\n",
+    "                    \"last_error\": last_error,\n",
+    "                })\n",
+    "    return rows\n",
+    "\n",
+    "\n",
+    "async def run_model_inference_each_task(label: str, seed: int = TEST_SEEDS[0], max_steps: int = 20):\n",
+    "    rows = []\n",
+    "    for scenario_id in SCENARIOS:\n",
+    "        print(f\"\\n=== {label}: {scenario_id} ===\")\n",
+    "        async with FlatmateEndpoint() as env:\n",
+    "            obs = await env.reset(scenario_id, seed=seed)\n",
+    "            total_reward = 0.0\n",
+    "            steps = 0\n",
+    "            parse_errors = 0\n",
+    "            for step_idx in range(max_steps):\n",
+    "                action, raw, error = model_action_or_error(obs)\n",
+    "                if action is None:\n",
+    "                    parse_errors += 1\n",
+    "                    print(f\"step={step_idx:02d} PARSE_ERROR={error}\")\n",
+    "                    print(\"raw=\", repr(raw[:300]))\n",
+    "                    total_reward -= 1.0\n",
+    "                    break\n",
+    "                print(f\"step={step_idx:02d} action={action}\")\n",
+    "                obs = await env.step(action)\n",
+    "                steps = step_idx + 1\n",
+    "                total_reward += float(obs.get(\"reward\") or obs.get(\"step_reward\") or 0.0)\n",
+    "                if obs.get(\"done\"):\n",
+    "                    break\n",
+    "            result = {\n",
+    "                \"policy\": label,\n",
+    "                \"scenario_id\": scenario_id,\n",
+    "                \"seed\": seed,\n",
+    "                \"total_reward\": total_reward,\n",
+    "                \"done\": bool(obs.get(\"done\")),\n",
+    "                \"bookings\": len(obs.get(\"booked_visits\", [])),\n",
+    "                \"violations\": len(obs.get(\"violations\", [])),\n",
+    "                \"steps\": steps,\n",
+    "                \"parse_errors\": parse_errors,\n",
+    "            }\n",
+    "            print(\"result=\", result)\n",
+    "            rows.append(result)\n",
+    "    return pd.DataFrame(rows)\n",
+    "\n",
+    "\n",
+    "active_model = base_model_for_eval\n",
+    "base_generation_sanity_df = await sanity_check_generations(\"base_model\")\n",
+    "base_per_task_inference_df = await run_model_inference_each_task(\"base_model\")\n",
+    "base_model_eval = await evaluate_model_policy(\"base_model\")\n",
+    "\n",
+    "active_model = loaded_model_for_eval\n",
+    "loaded_generation_sanity_df = await sanity_check_generations(\"sft_loaded\")\n",
+    "loaded_per_task_inference_df = await run_model_inference_each_task(\"sft_loaded\")\n",
+    "loaded_eval = await evaluate_model_policy(\"sft_loaded\")\n",
+    "\n",
+    "per_task_inference_df = pd.concat([base_per_task_inference_df, loaded_per_task_inference_df], ignore_index=True)\n",
+    "generation_sanity_df = pd.concat([base_generation_sanity_df, loaded_generation_sanity_df], ignore_index=True)\n",
+    "heuristic_eval = await evaluate_heuristic(\"heuristic\")\n",
+    "\n",
+    "eval_rows = heuristic_eval + base_model_eval + loaded_eval\n",
+    "eval_df = pd.DataFrame(eval_rows)\n",
+    "display(generation_sanity_df)\n",
+    "display(per_task_inference_df)\n",
+    "eval_df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e1e70c8f",
+   "metadata": {},
+   "source": [
+    "## Performance Comparison\n",
+    "\n",
+    "Compare heuristic rollout behavior against the trained SFT policy on the same scenarios and seeds."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8931930",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_policy_comparison(eval_df, title: str = \"Base vs SFT loaded-model comparison\"):\n",
+    "    if eval_df is None or eval_df.empty or \"policy\" not in eval_df.columns:\n",
+    "        print(\"eval_df is empty; run the evaluation cell first.\")\n",
+    "        return pd.DataFrame()\n",
+    "\n",
+    "    summary = (\n",
+    "        eval_df.groupby(\"policy\", as_index=True)\n",
+    "        .agg(\n",
+    "            avg_reward=(\"total_reward\", \"mean\"),\n",
+    "            completion_rate=(\"done\", \"mean\"),\n",
+    "            avg_bookings=(\"bookings\", \"mean\"),\n",
+    "            avg_violations=(\"violations\", \"mean\"),\n",
+    "            avg_steps=(\"steps\", \"mean\"),\n",
+    "            avg_parse_errors=(\"parse_errors\", \"mean\") if \"parse_errors\" in eval_df.columns else (\"steps\", \"size\"),\n",
+    "        )\n",
+    "        .sort_index()\n",
+    "    )\n",
+    "    plot_cols = [\"avg_reward\", \"completion_rate\", \"avg_bookings\", \"avg_violations\", \"avg_parse_errors\"]\n",
+    "    axes = summary[plot_cols].plot(\n",
+    "        kind=\"bar\",\n",
+    "        subplots=True,\n",
+    "        layout=(3, 2),\n",
+    "        figsize=(10, 9),\n",
+    "        legend=False,\n",
+    "        title=title,\n",
+    "    )\n",
+    "    for ax in axes.ravel():\n",
+    "        ax.grid(axis=\"y\", alpha=0.3)\n",
+    "        ax.set_xlabel(\"\")\n",
+    "    plt.tight_layout()\n",
+    "    plt.show()\n",
+    "    return summary\n",
+    "\n",
+    "comparison_summary = plot_policy_comparison(eval_df)\n",
+    "comparison_summary"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a9fd3807",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optional: upload the trained adapter/model to the Hub.\n",
+    "# from huggingface_hub import notebook_login\n",
+    "# notebook_login()\n",
+    "# trainer.push_to_hub(\"flatmate-rl-action-policy\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}