Upload folder using huggingface_hub

.pytest_cache/.gitignore

@@ -0,0 +1,2 @@
+# Created by pytest automatically.
+*

.pytest_cache/CACHEDIR.TAG

@@ -0,0 +1,4 @@
+Signature: 8a477f597d28d172789f06886806bc55
+# This file is a cache directory tag created by pytest.
+# For information about cache directory tags, see:
+#	https://bford.info/cachedir/spec.html

.pytest_cache/README.md

@@ -0,0 +1,8 @@
+# pytest cache directory #
+
+This directory contains data from the pytest's cache plugin,
+which provides the `--lf` and `--ff` options, as well as the `cache` fixture.
+
+**Do not** commit this to version control.
+
+See [the docs](https://docs.pytest.org/en/stable/how-to/cache.html) for more information.

.pytest_cache/v/cache/nodeids

@@ -0,0 +1,35 @@
+[
+  "tests/test_flatmate_rl.py::test_buyer_answers_diet_and_availability_when_broker_asks_for_both",
+  "tests/test_flatmate_rl.py::test_conflict_check_calendar_slots_exposes_pre_booked_and_available",
+  "tests/test_flatmate_rl.py::test_conflict_check_cannot_book_pre_booked_slot",
+  "tests/test_flatmate_rl.py::test_conflict_check_heuristic_books_only_available_slot",
+  "tests/test_flatmate_rl.py::test_heuristic_policy_progresses_after_confirmation_in_single_visit",
+  "tests/test_flatmate_rl.py::test_heuristic_policy_recovers_from_strict_eval_feedback",
+  "tests/test_flatmate_rl.py::test_hidden_flex_requires_alternative_slot_to_unlock_backup_availability",
+  "tests/test_flatmate_rl.py::test_legal_non_canonical_tool_after_store_can_continue",
+  "tests/test_flatmate_rl.py::test_multi_visit_scenario_books_two_visits",
+  "tests/test_flatmate_rl.py::test_negotiation_heuristic_confirms_deal_with_agreed_rent",
+  "tests/test_flatmate_rl.py::test_observation_surfaces_prerequisites_and_recent_tool_calls",
+  "tests/test_flatmate_rl.py::test_redundant_successful_tool_call_gets_small_penalty_without_termination",
+  "tests/test_flatmate_rl.py::test_reset_exposes_initial_buyer_message",
+  "tests/test_flatmate_rl.py::test_scenarios_are_self_consistent",
+  "tests/test_flatmate_rl.py::test_schema_valid_non_canonical_action_never_uses_legacy_flow_failure",
+  "tests/test_flatmate_rl.py::test_search_before_store_user_details_fails",
+  "tests/test_flatmate_rl.py::test_seeded_reset_varies_values_without_changing_episode_flow",
+  "tests/test_flatmate_rl.py::test_seller_followup_accepts_expected_tool_with_different_arguments",
+  "tests/test_flatmate_rl.py::test_seller_followup_accepts_paraphrased_assistant_message",
+  "tests/test_flatmate_rl.py::test_seller_followup_match_tools_infer_dynamic_post_when_args_are_loose",
+  "tests/test_flatmate_rl.py::test_seller_followup_non_canonical_tool_order_gets_small_penalty_without_termination",
+  "tests/test_flatmate_rl.py::test_seller_followup_scenario_schedules_dynamic_visit",
+  "tests/test_flatmate_rl.py::test_seller_followup_search_returns_no_visit_compatible_current_posts",
+  "tests/test_flatmate_rl.py::test_single_visit_scenario_books_one_visit",
+  "tests/test_flatmate_rl.py::test_store_user_details_does_not_return_expected_answers_payload",
+  "tests/test_flatmate_rl.py::test_strict_eval_mode_hides_scenario_metadata_and_reward",
+  "tests/test_inference_parsing.py::test_legacy_parse_can_coerce_tool_name_in_action_type",
+  "tests/test_inference_parsing.py::test_malformed_action_feedback_is_recoverable",
+  "tests/test_inference_parsing.py::test_model_call_error_does_not_fallback_to_heuristic",
+  "tests/test_inference_parsing.py::test_strict_parse_rejects_tool_name_in_action_type",
+  "tests/test_inference_parsing.py::test_strict_parse_reports_json_error",
+  "tests/test_inference_parsing.py::test_user_prompt_renders_prerequisites_and_recent_tools",
+  "tests/test_reward_regression.py::test_heuristic_reward_regression"
+]

server/episode.py

@@ -11,10 +11,12 @@ from typing import Any
 try:
     from ..models import FlatmateRlAction, FlatmateRlObservation, FlatmateRlState
     from .heuristic_policy import expected_policy_action
+    from .scenario_variants import apply_seed_variant
     from .scenarios import POSTS, SCENARIOS
 except ImportError:
     from models import FlatmateRlAction, FlatmateRlObservation, FlatmateRlState
     from server.heuristic_policy import expected_policy_action
+    from server.scenario_variants import apply_seed_variant
     from server.scenarios import POSTS, SCENARIOS
 
 
@@ -118,10 +120,11 @@ class FlatmateEpisode:
         self._post_arrivals_fired: set[int] = set()
         self._available_post_ids: list[str] = []
 
-    def reset(self, scenario_id: str | None = None) -> FlatmateRlObservation:
+    def reset(self, scenario_id: str | None = None, seed: int | None = None) -> FlatmateRlObservation:
         selected = scenario_id or "task_visit_single"
-        self._scenario = deepcopy(SCENARIOS[selected])
-        self._posts = {post_id: deepcopy(POSTS[post_id]) for post_id in self._scenario["task_post_ids"]}
+        base_scenario = deepcopy(SCENARIOS[selected])
+        base_posts = {post_id: deepcopy(POSTS[post_id]) for post_id in base_scenario["task_post_ids"]}
+        self._scenario, self._posts = apply_seed_variant(base_scenario, base_posts, seed)
         self._tool_results = []
         self._tool_trace = []
         self._history = []

server/flatmate_rl_environment.py

@@ -27,9 +27,8 @@ class FlatmateRlEnvironment(Environment):
         self._episode = FlatmateEpisode(strict_eval_mode=strict_eval_mode)
 
     def reset(self, scenario_id: str | None = None, seed: int | None = None) -> FlatmateRlObservation:
-        del seed
         self._state = FlatmateRlState(episode_id=str(uuid4()), step_count=0)
-        observation = self._episode.reset(scenario_id=scenario_id)
+        observation = self._episode.reset(scenario_id=scenario_id, seed=seed)
         self._sync_state(observation)
         return observation
 

server/scenario_variants.py

@@ -0,0 +1,133 @@
+"""Seeded value variants for Flatmate RL scenarios.
+
+Variants intentionally preserve episode flow: same task id, post ids, required
+tools, required bookings, feasible slots, and phase transitions. Only safe
+surface values are shifted so train/test episodes can differ without changing
+the canonical solution structure.
+"""
+
+from __future__ import annotations
+
+import random
+from copy import deepcopy
+from typing import Any
+
+
+OCCUPATIONS = [
+    "software engineer at a startup",
+    "backend engineer at a fintech company",
+    "data engineer at a healthtech company",
+    "platform engineer at a SaaS company",
+]
+
+RENT_DELTAS = [-1200, -800, -500, 0, 500, 800, 1200]
+
+
+def _format_rs(amount: int) -> str:
+    return f"{amount:,}"
+
+
+def _buyer_message(scenario: dict[str, Any]) -> str:
+    buyer = scenario["buyer_profile"]
+    areas = " or ".join(buyer["areas"])
+    budget = _format_rs(int(buyer["budget_max"]))
+    occupation = buyer["occupation"]
+    availability = " or ".join(buyer["visit_availability"])
+    task_id = scenario["task_id"]
+
+    if task_id == "task_visit_multi":
+        return (
+            "Hi, I want to line up visits for at least two good flatmate-share options before deciding. "
+            f"My budget is Rs. {budget}, I'm focused on {areas}, and I work in Goregaon East as a {occupation}."
+        )
+    if task_id == "task_visit_single_seller_followup":
+        return (
+            f"Hi, I'm looking for a flatmate-share in {areas}. My budget is Rs. {budget}, "
+            f"I work in Goregaon East as a {occupation}, and {availability} are the only times I can visit."
+        )
+    if "visit_availability" in buyer.get("initial_disclosure_fields", []):
+        return (
+            f"Hi, I'm looking for a flatmate-share around {areas}. My budget is Rs. {budget}, "
+            f"I work in Goregaon East as a {occupation}, and {availability} is the slot I can do right now."
+        )
+    return (
+        "Hi, I'm looking for a flatmate-share near Goregaon East. "
+        f"My budget is up to Rs. {budget} and I'm mainly considering {areas} "
+        f"because I work as a {occupation}."
+    )
+
+
+def _seller_message(scenario: dict[str, Any]) -> str:
+    seller = scenario.get("seller_profile")
+    if not seller:
+        return scenario.get("seller_initial_message", "")
+    return (
+        f"Hi, I want help listing a new flatmate-share opening in {seller['area']}. "
+        f"The rent is around Rs. {_format_rs(int(seller['rent']))} for a {seller['listing_type']}. "
+        "I can tell you more about the listing and available visit times."
+    )
+
+
+def _shift_amount(value: Any, delta: int) -> Any:
+    if value is None:
+        return None
+    return max(1, int(value) + delta)
+
+
+def apply_seed_variant(
+    scenario: dict[str, Any],
+    posts: dict[str, dict[str, Any]],
+    seed: int | None,
+) -> tuple[dict[str, Any], dict[str, dict[str, Any]]]:
+    """Return seeded copies of a scenario and its posts.
+
+    The same seed always produces the same value variant for a scenario. A
+    missing seed returns the original values, preserving existing tests and
+    default behavior.
+    """
+
+    variant_scenario = deepcopy(scenario)
+    variant_posts = deepcopy(posts)
+    if seed is None:
+        return variant_scenario, variant_posts
+
+    rng = random.Random(f"{variant_scenario['task_id']}:{seed}")
+    rent_delta = rng.choice(RENT_DELTAS)
+    occupation = rng.choice(OCCUPATIONS)
+
+    buyer = variant_scenario["buyer_profile"]
+    buyer["budget_max"] = _shift_amount(buyer["budget_max"], rent_delta)
+    if buyer.get("hidden_budget_ceiling") is not None:
+        buyer["hidden_budget_ceiling"] = _shift_amount(buyer["hidden_budget_ceiling"], rent_delta)
+    buyer["occupation"] = occupation
+
+    expected_answers = variant_scenario["scenario_creation_config"].get("expected_answers", {})
+    if "budget_max" in expected_answers:
+        expected_answers["budget_max"] = buyer["budget_max"]
+    if "occupation" in expected_answers:
+        expected_answers["occupation"] = buyer["occupation"]
+
+    negotiation_config = variant_scenario["scenario_creation_config"].get("negotiation_config", {})
+    for key in ("buyer_ceiling", "seller_floor"):
+        if key in negotiation_config:
+            negotiation_config[key] = _shift_amount(negotiation_config[key], rent_delta)
+
+    for post in variant_posts.values():
+        post["rent"] = _shift_amount(post["rent"], rent_delta)
+
+    seller = variant_scenario.get("seller_profile")
+    if seller:
+        seller["rent"] = _shift_amount(seller["rent"], rent_delta)
+        followup = variant_scenario["scenario_creation_config"].get("followup_seller_expected_answers", {})
+        if "rent" in followup:
+            followup["rent"] = seller["rent"]
+
+    variant_scenario["initial_user_message"] = _buyer_message(variant_scenario)
+    if seller:
+        variant_scenario["seller_initial_message"] = _seller_message(variant_scenario)
+    variant_scenario["scenario_creation_config"]["variant"] = {
+        "seed": seed,
+        "rent_delta": rent_delta,
+        "occupation": occupation,
+    }
+    return variant_scenario, variant_posts

tests/test_flatmate_rl.py

@@ -84,6 +84,25 @@ def test_reset_exposes_initial_buyer_message() -> None:
     assert observation.remaining_required_fields == ["diet", "visit_availability"]
 
 
+def test_seeded_reset_varies_values_without_changing_episode_flow() -> None:
+    default_env = FlatmateRlEnvironment()
+    default_obs = default_env.reset(scenario_id="task_visit_single")
+
+    seeded_env = FlatmateRlEnvironment()
+    seeded_obs = seeded_env.reset(scenario_id="task_visit_single", seed=123)
+    seeded_episode = seeded_env._episode  # type: ignore[attr-defined]
+
+    assert seeded_obs.scenario_id == default_obs.scenario_id
+    assert seeded_obs.remaining_required_fields == default_obs.remaining_required_fields
+    assert seeded_episode._scenario["task_post_ids"] == SCENARIOS["task_visit_single"]["task_post_ids"]
+    assert seeded_episode._scenario["ground_truth"] == SCENARIOS["task_visit_single"]["ground_truth"]
+    assert seeded_episode._scenario["buyer_profile"]["budget_max"] != SCENARIOS["task_visit_single"]["buyer_profile"]["budget_max"]
+    assert (
+        seeded_episode._scenario["scenario_creation_config"]["expected_answers"]["budget_max"]
+        == seeded_episode._scenario["buyer_profile"]["budget_max"]
+    )
+
+
 def test_search_before_store_user_details_fails() -> None:
     env = FlatmateRlEnvironment()
     env.reset(scenario_id="task_visit_single")

train_flatmate_rl_grpo.ipynb

@@ -0,0 +1,838 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f1f7a30f",
+   "metadata": {},
+   "source": [
+    "# Train Flatmate RL with TRL GRPO\n",
+    "\n",
+    "This notebook uses the live Hugging Face Space endpoint as the reward source for GRPO. It first collects prompt states from websocket rollouts, then trains a model to generate one JSON action for each observation. The main reward function resets the endpoint, replays the stored action prefix, applies the model action, and scores the resulting Flatmate RL transition.\n",
+    "\n",
+    "Space endpoint: `https://huggingface.co/spaces/kushalExplores/flatmate_rl`"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "adabfca9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Restart the kernel after this cell if your notebook runtime asks you to.\n",
+    "%pip install -q \"trl>=0.26.2\" \"transformers>=4.57.0\" accelerate datasets peft websockets huggingface_hub matplotlib pandas"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b79df9da",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations\n",
+    "\n",
+    "import asyncio\n",
+    "import json\n",
+    "import random\n",
+    "import threading\n",
+    "from dataclasses import dataclass\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": "ac8bdbfc",
+   "metadata": {},
+   "source": [
+    "## Websocket Environment Client\n",
+    "\n",
+    "OpenEnv keeps episode state on `/ws`. The plain HTTP `/reset` and `/step` endpoints create fresh environment instances, so GRPO reward replay uses websocket sessions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "bf61d6d4",
+   "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(\n",
+    "            self.ws_url,\n",
+    "            open_timeout=self.timeout_s,\n",
+    "            ping_timeout=self.timeout_s,\n",
+    "        )\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": "d6401b21",
+   "metadata": {},
+   "source": [
+    "## Prompt States\n",
+    "\n",
+    "GRPO needs prompts. These prompts are endpoint observations collected from heuristic rollouts. Each row also stores the action prefix needed to reconstruct that exact state during reward scoring."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c2ebe016",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def trace_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 heuristic_action(obs: dict[str, Any]) -> dict[str, Any] | None:\n",
+    "    tools = trace_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",
+    "    post_ids = [\"post_031\", \"post_052\"] if scenario_id == \"task_visit_multi\" else [\"post_031\"]\n",
+    "    if \"search_posts\" not in tools:\n",
+    "        return {\"action_type\": \"tool_call\", \"tool_name\": \"search_posts\", \"tool_arguments\": {}}\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",
+    "    return None\n",
+    "\n",
+    "def compact_observation(obs: dict[str, Any]) -> dict[str, Any]:\n",
+    "    return {\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",
+    "\n",
+    "def prompt_from_observation(obs: dict[str, Any]) -> str:\n",
+    "    return (\n",
+    "        \"You are a broker policy for the Flatmate RL environment. \"\n",
+    "        \"Return exactly one JSON action and no extra text.\\n\\n\"\n",
+    "        \"Valid action shapes:\\n\"\n",
+    "        \"{\\\"action_type\\\":\\\"assistant_message\\\",\\\"assistant_message\\\":\\\"...\\\"}\\n\"\n",
+    "        \"{\\\"action_type\\\":\\\"tool_call\\\",\\\"tool_name\\\":\\\"...\\\",\\\"tool_arguments\\\":{...}}\\n\\n\"\n",
+    "        f\"Observation:\\n{json.dumps(compact_observation(obs), ensure_ascii=False, sort_keys=True)}\\n\\n\"\n",
+    "        \"Action:\\n\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d29c9dc5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "@dataclass\n",
+    "class PromptCollectionConfig:\n",
+    "    episodes: int = 12\n",
+    "    max_steps: int = 14\n",
+    "    seed: int = 11\n",
+    "\n",
+    "async def collect_prompt_rows(config: PromptCollectionConfig = PromptCollectionConfig()) -> list[dict[str, Any]]:\n",
+    "    rng = random.Random(config.seed)\n",
+    "    rows: list[dict[str, Any]] = []\n",
+    "    for episode_idx in range(config.episodes):\n",
+    "        scenario_id = rng.choice(SCENARIOS)\n",
+    "        prefix_actions: list[dict[str, Any]] = []\n",
+    "        async with FlatmateEndpoint() as env:\n",
+    "            obs = await env.reset(scenario_id, seed=config.seed + episode_idx)\n",
+    "            for step_idx in range(config.max_steps):\n",
+    "                action = heuristic_action(obs)\n",
+    "                if action is None or obs.get(\"done\"):\n",
+    "                    break\n",
+    "                rows.append({\n",
+    "                    \"prompt\": prompt_from_observation(obs),\n",
+    "                    \"scenario_id\": scenario_id,\n",
+    "                    \"seed\": config.seed + episode_idx,\n",
+    "                    \"prefix_actions_json\": json.dumps(prefix_actions, ensure_ascii=False),\n",
+    "                    \"reference_action_json\": json.dumps(action, ensure_ascii=False, sort_keys=True),\n",
+    "                })\n",
+    "                obs = await env.step(action)\n",
+    "                prefix_actions.append(action)\n",
+    "                if obs.get(\"done\"):\n",
+    "                    break\n",
+    "        print(f\"episode={episode_idx:02d} scenario={scenario_id} total_rows={len(rows)}\")\n",
+    "    return rows\n",
+    "\n",
+    "rows = await collect_prompt_rows(PromptCollectionConfig(episodes=12, max_steps=14))\n",
+    "train_dataset = Dataset.from_list(rows)\n",
+    "train_dataset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a10b4cd1",
+   "metadata": {},
+   "source": [
+    "## GRPO Rewards\n",
+    "\n",
+    "`json_format_reward` is cheap and runs for every completion. `endpoint_transition_reward` is slower because it calls the live Space: it replays the prefix actions, sends the generated action, and returns a reward based on environment reward, validity, violations, and bookings."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e7bdfc62",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def completion_text(completion: Any) -> str:\n",
+    "    if isinstance(completion, str):\n",
+    "        return completion\n",
+    "    if isinstance(completion, list) and completion and isinstance(completion[0], dict):\n",
+    "        return str(completion[0].get(\"content\", \"\"))\n",
+    "    return str(completion)\n",
+    "\n",
+    "def parse_json_action(text: str) -> dict[str, Any]:\n",
+    "    text = completion_text(text).strip()\n",
+    "    start = text.find(\"{\")\n",
+    "    end = text.rfind(\"}\")\n",
+    "    if start == -1 or end == -1 or end <= start:\n",
+    "        raise ValueError(\"completion does not contain a JSON object\")\n",
+    "    action = json.loads(text[start : end + 1])\n",
+    "    if action.get(\"action_type\") == \"assistant_message\":\n",
+    "        if not str(action.get(\"assistant_message\", \"\")).strip():\n",
+    "            raise ValueError(\"assistant_message action needs assistant_message\")\n",
+    "        return {\n",
+    "            \"action_type\": \"assistant_message\",\n",
+    "            \"assistant_message\": str(action[\"assistant_message\"]),\n",
+    "        }\n",
+    "    if action.get(\"action_type\") == \"tool_call\":\n",
+    "        if not str(action.get(\"tool_name\", \"\")).strip():\n",
+    "            raise ValueError(\"tool_call action needs tool_name\")\n",
+    "        args = action.get(\"tool_arguments\", {})\n",
+    "        if not isinstance(args, dict):\n",
+    "            raise ValueError(\"tool_arguments must be an object\")\n",
+    "        return {\n",
+    "            \"action_type\": \"tool_call\",\n",
+    "            \"tool_name\": str(action[\"tool_name\"]),\n",
+    "            \"tool_arguments\": args,\n",
+    "        }\n",
+    "    raise ValueError(\"action_type must be assistant_message or tool_call\")\n",
+    "\n",
+    "def json_format_reward(completions, **kwargs) -> list[float]:\n",
+    "    rewards = []\n",
+    "    for completion in completions:\n",
+    "        try:\n",
+    "            parse_json_action(completion_text(completion))\n",
+    "            rewards.append(0.25)\n",
+    "        except Exception:\n",
+    "            rewards.append(-1.0)\n",
+    "    return rewards\n",
+    "\n",
+    "async def score_one_completion(\n",
+    "    completion: Any,\n",
+    "    scenario_id: str,\n",
+    "    seed: int,\n",
+    "    prefix_actions_json: str,\n",
+    ") -> float:\n",
+    "    try:\n",
+    "        action = parse_json_action(completion_text(completion))\n",
+    "        prefix_actions = json.loads(prefix_actions_json)\n",
+    "    except Exception:\n",
+    "        return -1.0\n",
+    "\n",
+    "    try:\n",
+    "        async with FlatmateEndpoint() as env:\n",
+    "            obs = await env.reset(scenario_id, seed=int(seed))\n",
+    "            for prefix_action in prefix_actions:\n",
+    "                obs = await env.step(prefix_action)\n",
+    "                if obs.get(\"done\"):\n",
+    "                    return -0.5\n",
+    "\n",
+    "            before_violations = len(obs.get(\"violations\", []))\n",
+    "            before_bookings = len(obs.get(\"booked_visits\", []))\n",
+    "            obs = await env.step(action)\n",
+    "\n",
+    "        reward = float(obs.get(\"reward\") or obs.get(\"step_reward\") or 0.0)\n",
+    "        reward += 0.15\n",
+    "        if len(obs.get(\"violations\", [])) > before_violations:\n",
+    "            reward -= 0.75\n",
+    "        if len(obs.get(\"booked_visits\", [])) > before_bookings:\n",
+    "            reward += 1.0\n",
+    "        if obs.get(\"done\"):\n",
+    "            reward += 0.5\n",
+    "        return float(max(-2.0, min(2.0, reward)))\n",
+    "    except Exception as exc:\n",
+    "        print(\"endpoint reward error:\", repr(exc))\n",
+    "        return -1.0\n",
+    "\n",
+    "async def score_completion_batch(completions, scenario_id, seed, prefix_actions_json) -> list[float]:\n",
+    "    tasks = [\n",
+    "        score_one_completion(c, s, int(sd), p)\n",
+    "        for c, s, sd, p in zip(completions, scenario_id, seed, prefix_actions_json)\n",
+    "    ]\n",
+    "    return list(await asyncio.gather(*tasks))\n",
+    "\n",
+    "def run_async_blocking(coro):\n",
+    "    try:\n",
+    "        asyncio.get_running_loop()\n",
+    "    except RuntimeError:\n",
+    "        return asyncio.run(coro)\n",
+    "\n",
+    "    result: dict[str, Any] = {}\n",
+    "    def runner():\n",
+    "        try:\n",
+    "            result[\"value\"] = asyncio.run(coro)\n",
+    "        except Exception as exc:\n",
+    "            result[\"error\"] = exc\n",
+    "\n",
+    "    thread = threading.Thread(target=runner, daemon=True)\n",
+    "    thread.start()\n",
+    "    thread.join()\n",
+    "    if \"error\" in result:\n",
+    "        raise result[\"error\"]\n",
+    "    return result[\"value\"]\n",
+    "\n",
+    "def endpoint_transition_reward(completions, scenario_id, seed, prefix_actions_json, **kwargs) -> list[float]:\n",
+    "    return run_async_blocking(score_completion_batch(completions, scenario_id, seed, prefix_actions_json))\n",
+    "\n",
+    "# Quick reward smoke test with known reference actions from the collected dataset.\n",
+    "sample = train_dataset.select(range(min(2, len(train_dataset))))\n",
+    "endpoint_transition_reward(\n",
+    "    completions=sample[\"reference_action_json\"],\n",
+    "    scenario_id=sample[\"scenario_id\"],\n",
+    "    seed=sample[\"seed\"],\n",
+    "    prefix_actions_json=sample[\"prefix_actions_json\"],\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d286bdc7",
+   "metadata": {},
+   "source": [
+    "## Train with GRPO\n",
+    "\n",
+    "The endpoint reward is network-bound, so keep `num_generations`, dataset size, and max steps small until the loop is stable. Increase them once you see valid JSON actions and non-negative endpoint rewards."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6323589",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import inspect\n",
+    "\n",
+    "from peft import LoraConfig\n",
+    "from trl import GRPOConfig, GRPOTrainer\n",
+    "\n",
+    "MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
+    "OUTPUT_DIR = \"flatmate-rl-grpo-policy\"\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",
+    "def make_grpo_config(**kwargs):\n",
+    "    valid = set(inspect.signature(GRPOConfig.__init__).parameters)\n",
+    "    filtered = {key: value for key, value in kwargs.items() if key in valid}\n",
+    "    skipped = sorted(set(kwargs) - set(filtered))\n",
+    "    if skipped:\n",
+    "        print(\"Skipping unsupported GRPOConfig fields for this TRL version:\", skipped)\n",
+    "    return GRPOConfig(**filtered)\n",
+    "\n",
+    "training_args = make_grpo_config(\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    learning_rate=1e-5,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    num_generations=4,\n",
+    "    max_prompt_length=1536,\n",
+    "    max_completion_length=160,\n",
+    "    max_steps=30,\n",
+    "    logging_steps=1,\n",
+    "    save_steps=15,\n",
+    "    save_total_limit=2,\n",
+    "    report_to=\"none\",\n",
+    ")\n",
+    "\n",
+    "trainer = GRPOTrainer(\n",
+    "    model=MODEL_NAME,\n",
+    "    reward_funcs=[json_format_reward, endpoint_transition_reward],\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_dataset,\n",
+    "    peft_config=peft_config,\n",
+    ")\n",
+    "\n",
+    "train_result = trainer.train()\n",
+    "train_log_history = trainer.state.log_history\n",
+    "trainer.save_model(OUTPUT_DIR)\n",
+    "train_result"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "da41a231",
+   "metadata": {},
+   "source": [
+    "## Training Log\n",
+    "\n",
+    "Plot logged GRPO reward and loss curves over optimizer steps."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a6c8071d",
+   "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 = \"GRPO training log\"):\n",
+    "    df = pd.DataFrame(log_history)\n",
+    "    if df.empty:\n",
+    "        print(\"No trainer log rows found yet.\")\n",
+    "        return df\n",
+    "    metric_cols = [col for col in [\"loss\", \"reward\", \"rewards/json_format_reward\", \"rewards/endpoint_transition_reward\", \"kl\"] if col in df.columns]\n",
+    "    if not metric_cols:\n",
+    "        metric_cols = [col for col in df.columns if \"reward\" in col or col in {\"loss\", \"kl\"}]\n",
+    "    if not metric_cols or \"step\" not in df.columns:\n",
+    "        print(\"No plottable step metrics found. Available columns:\", list(df.columns))\n",
+    "        return df\n",
+    "    axes = df.dropna(subset=[\"step\"]).plot(\n",
+    "        x=\"step\",\n",
+    "        y=metric_cols,\n",
+    "        marker=\"o\",\n",
+    "        figsize=(9, 4),\n",
+    "        title=title,\n",
+    "    )\n",
+    "    axes.set_xlabel(\"optimizer step\")\n",
+    "    axes.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()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c013682e",
+   "metadata": {},
+   "source": [
+    "## Evaluate One Episode"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4cc3bc3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from peft import AutoPeftModelForCausalLM\n",
+    "from transformers import AutoModelForCausalLM, AutoTokenizer\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",
+    "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 GRPO model from {OUTPUT_DIR}\")\n",
+    "\n",
+    "generation_stats = {\"fallbacks\": 0, \"parse_errors\": []}\n",
+    "\n",
+    "\n",
+    "def safe_model_action(obs: dict[str, Any], log_bad: bool = False) -> dict[str, Any]:\n",
+    "    # Prefix the completion with \"{\" so the model only has to finish a JSON object.\n",
+    "    prompt = prompt_from_observation(obs) + \"{\"\n",
+    "    inputs = tokenizer(prompt, return_tensors=\"pt\", truncation=True, max_length=1536).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",
+    "    completion_tail = tokenizer.decode(output[0][inputs[\"input_ids\"].shape[-1]:], skip_special_tokens=True)\n",
+    "    completion = \"{\" + completion_tail\n",
+    "    try:\n",
+    "        return parse_json_action(completion)\n",
+    "    except Exception as exc:\n",
+    "        generation_stats[\"fallbacks\"] += 1\n",
+    "        if len(generation_stats[\"parse_errors\"]) < 5:\n",
+    "            generation_stats[\"parse_errors\"].append({\"error\": str(exc), \"completion\": completion[:300]})\n",
+    "        if log_bad:\n",
+    "            print(f\"bad generation, using fallback: {exc}\")\n",
+    "            print(\"raw completion:\", repr(completion[:300]))\n",
+    "        fallback = heuristic_action(obs)\n",
+    "        if fallback is None:\n",
+    "            fallback = {\"action_type\": \"assistant_message\", \"assistant_message\": \"Could you confirm the details needed for scheduling?\"}\n",
+    "        return fallback\n",
+    "\n",
+    "def heuristic_policy(obs: dict[str, Any]) -> dict[str, Any]:\n",
+    "    action = heuristic_action(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",
+    "async def evaluate_policy(policy_fn, label: str, scenarios=SCENARIOS, seeds=(123, 124), max_steps: int = 20, verbose: bool = False):\n",
+    "    rows = []\n",
+    "    for scenario_id in scenarios:\n",
+    "        for seed in seeds:\n",
+    "            before_fallbacks = generation_stats[\"fallbacks\"] if \"generation_stats\" in globals() else 0\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 = policy_fn(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",
+    "                after_fallbacks = generation_stats[\"fallbacks\"] if \"generation_stats\" in globals() else before_fallbacks\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",
+    "                    \"fallbacks\": after_fallbacks - before_fallbacks,\n",
+    "                })\n",
+    "    return rows\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\", truncation=True, max_length=1536).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",
+    "async def sanity_check_generations(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=321)\n",
+    "            raw = raw_generate_action_text(obs)\n",
+    "            try:\n",
+    "                parsed = parse_json_action(raw)\n",
+    "                ok = True\n",
+    "            except Exception as exc:\n",
+    "                parsed = str(exc)\n",
+    "                ok = False\n",
+    "            rows.append({\"scenario_id\": scenario_id, \"json_ok\": ok, \"raw\": raw[:240], \"parsed_or_error\": parsed})\n",
+    "    return pd.DataFrame(rows)\n",
+    "\n",
+    "async def run_inference_each_task(policy_fn, label: str, seed: int = 321, 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",
+    "            before_fallbacks = generation_stats[\"fallbacks\"]\n",
+    "            for step_idx in range(max_steps):\n",
+    "                action = policy_fn(obs)\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",
+    "                \"fallbacks\": generation_stats[\"fallbacks\"] - before_fallbacks,\n",
+    "            }\n",
+    "            print(\"result=\", result)\n",
+    "            rows.append(result)\n",
+    "    return pd.DataFrame(rows)\n",
+    "\n",
+    "generation_stats = {\"fallbacks\": 0, \"parse_errors\": []}\n",
+    "active_model = base_model_for_eval\n",
+    "base_generation_sanity_df = await sanity_check_generations()\n",
+    "base_generation_sanity_df[\"model\"] = \"base\"\n",
+    "base_per_task_inference_df = await run_inference_each_task(safe_model_action, \"base_model\")\n",
+    "base_model_eval = await evaluate_policy(safe_model_action, \"base_model\")\n",
+    "base_stats = dict(generation_stats)\n",
+    "\n",
+    "active_model = loaded_model_for_eval\n",
+    "generation_stats = {\"fallbacks\": 0, \"parse_errors\": []}\n",
+    "loaded_generation_sanity_df = await sanity_check_generations()\n",
+    "loaded_generation_sanity_df[\"model\"] = \"grpo_loaded\"\n",
+    "loaded_per_task_inference_df = await run_inference_each_task(safe_model_action, \"grpo_loaded\")\n",
+    "per_task_inference_df = pd.concat([base_per_task_inference_df, loaded_per_task_inference_df], ignore_index=True)\n",
+    "loaded_stats = dict(generation_stats)\n",
+    "\n",
+    "generation_sanity_df = pd.concat([base_generation_sanity_df, loaded_generation_sanity_df], ignore_index=True)\n",
+    "print(\"base_generation_stats\", base_stats)\n",
+    "print(\"loaded_generation_stats\", loaded_stats)\n",
+    "if loaded_stats[\"fallbacks\"] > 0:\n",
+    "    print(\"WARNING: loaded fine-tuned model produced malformed JSON and used fallback.\")\n",
+    "\n",
+    "heuristic_eval = await evaluate_policy(heuristic_policy, \"heuristic\")\n",
+    "active_model = loaded_model_for_eval\n",
+    "grpo_eval = await evaluate_policy(safe_model_action, \"grpo_loaded\")\n",
+    "eval_rows = heuristic_eval + base_model_eval + grpo_eval\n",
+    "eval_df = pd.DataFrame(eval_rows)\n",
+    "eval_df"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "644fd566",
+   "metadata": {},
+   "source": [
+    "## Performance Comparison\n",
+    "\n",
+    "Compare heuristic rollout behavior against the trained GRPO policy on the same scenarios and seeds."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0162e4ad",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def plot_policy_comparison(eval_df, title: str = \"Base vs GRPO 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 because loaded-model generation used fallback or evaluation was skipped.\")\n",
+    "        if \"per_task_inference_df\" in globals() and not per_task_inference_df.empty:\n",
+    "            print(\"Plotting per_task_inference_df instead. This is loaded-model rollout only, not heuristic-vs-model comparison.\")\n",
+    "            display(per_task_inference_df)\n",
+    "            ax = per_task_inference_df.set_index(\"scenario_id\")[[\"total_reward\", \"bookings\", \"violations\", \"fallbacks\"]].plot(\n",
+    "                kind=\"bar\",\n",
+    "                subplots=True,\n",
+    "                layout=(2, 2),\n",
+    "                figsize=(10, 7),\n",
+    "                legend=False,\n",
+    "                title=\"GRPO loaded-model per-task inference\",\n",
+    "            )\n",
+    "            for axis in ax.ravel():\n",
+    "                axis.grid(axis=\"y\", alpha=0.3)\n",
+    "                axis.set_xlabel(\"\")\n",
+    "            plt.tight_layout()\n",
+    "            plt.show()\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_fallbacks=(\"fallbacks\", \"mean\") if \"fallbacks\" in eval_df.columns else (\"steps\", \"size\"),\n",
+    "        )\n",
+    "        .sort_index()\n",
+    "    )\n",
+    "    plot_cols = [\"avg_reward\", \"completion_rate\", \"avg_bookings\", \"avg_violations\"]\n",
+    "    if \"avg_fallbacks\" in summary.columns:\n",
+    "        plot_cols.append(\"avg_fallbacks\")\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": "ee2e4887",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Optional Hub upload after training.\n",
+    "# from huggingface_hub import notebook_login\n",
+    "# notebook_login()\n",
+    "# trainer.push_to_hub(\"flatmate-rl-grpo-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
+}

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
+}