fix(training): reward signal flow + new v2 notebook

- rewards.py: add defensive _extract warning when kwargs are empty;
add combined_reward() as single-function fallback; add COMBINED_REWARD_FUNC list
- rollout.py: print per-episode reward breakdown; compute combined reward inline
- train_v2.ipynb: new notebook addressing all v1 issues:
* pins TRL 1.2.0 stable (no git install)
* removes PatchFastRL (incompatible with TRL 1.2.0)
* single 'reward' key from rollout (bypasses kwargs-flow bug)
* auto-detects T4 vs A100, adapts grad_accum/num_gen/max_comp
* smoke test cell before training

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.gitignore

@@ -12,3 +12,4 @@ pdf_pages/
 dist/
 build/
 .pytest_cache/
+.vscode/

training/rewards.py

@@ -16,8 +16,21 @@ Reward hacking protection (point 8 of hackathon guide):
 
 To change weights, edit the WEIGHT_* constants below.
 Equal-weight Option B: set all five to 0.20.
+
+Architecture note — two reward modes:
+  1. Split mode (ALL_REWARD_FUNCS): five separate functions, each reading one
+     kwargs key from the rollout dict. TRL sums them. Requires TRL to pass
+     rollout extra keys as kwargs to reward_funcs (works in TRL 0.17–0.24 +
+     Unsloth patch).
+  2. Combined mode (COMBINED_REWARD_FUNC): single function that reads all five
+     kwargs keys and returns the weighted sum. Use this if split mode shows
+     all-zero rewards (indicates TRL is not passing kwargs, e.g. version
+     mismatch). Switch by replacing reward_funcs=ALL_REWARD_FUNCS with
+     reward_funcs=COMBINED_REWARD_FUNC in GRPOTrainer.
 """
 
+import warnings
+
 WEIGHT_FINAL_SCORE = 0.45
 WEIGHT_NO_WRONG_CHANNELS = 0.15
 WEIGHT_VALID_JSON = 0.15
@@ -29,10 +42,32 @@ assert abs(
     + WEIGHT_READ_RUNBOOK + WEIGHT_EFFICIENCY - 1.0
 ) < 1e-9, "Reward weights must sum to 1.0"
 
+_EXPECTED_KEYS = frozenset([
+    "final_score_reward", "no_wrong_channels_reward",
+    "valid_json_reward", "read_runbook_reward", "efficiency_reward",
+])
+_warned_empty = False
+
 
 def _extract(kwargs: dict, key: str, n: int) -> list[float]:
+    global _warned_empty
     rewards = kwargs.get(key, [])
-    return [float(r) for r in rewards] if rewards else [0.0] * n
+    if not rewards:
+        if not _warned_empty:
+            present = set(kwargs.keys()) & _EXPECTED_KEYS
+            missing = _EXPECTED_KEYS - set(kwargs.keys())
+            warnings.warn(
+                f"[rewards] kwargs missing rollout reward keys — all rewards will be 0!\n"
+                f"  present: {present or 'none'}\n"
+                f"  missing: {missing}\n"
+                f"  all kwargs keys: {list(kwargs.keys())}\n"
+                f"  CAUSE: TRL version mismatch or rollout_func not returning these keys.\n"
+                f"  FIX: switch to reward_funcs=COMBINED_REWARD_FUNC (see rewards.py).",
+                stacklevel=3,
+            )
+            _warned_empty = True
+        return [0.0] * n
+    return [float(r) for r in rewards]
 
 
 def reward_final_score(completions, **kwargs) -> list[float]:
@@ -79,7 +114,29 @@ def reward_efficiency(completions, **kwargs) -> list[float]:
     return [r * WEIGHT_EFFICIENCY for r in raw]
 
 
-# Ordered list for GRPOTrainer(reward_funcs=...)
+def combined_reward(completions, **kwargs) -> list[float]:
+    """Single combined reward — fallback when TRL does not pass split kwargs.
+
+    Reads all five rollout reward keys and returns their weighted sum.
+    Use with: reward_funcs=[combined_reward]
+    """
+    n = len(completions)
+    final    = _extract(kwargs, "final_score_reward",       n)
+    no_wrong = _extract(kwargs, "no_wrong_channels_reward", n)
+    vj       = _extract(kwargs, "valid_json_reward",        n)
+    rb       = _extract(kwargs, "read_runbook_reward",      n)
+    eff      = _extract(kwargs, "efficiency_reward",        n)
+    return [
+        f * WEIGHT_FINAL_SCORE
+        + nw * WEIGHT_NO_WRONG_CHANNELS
+        + v  * WEIGHT_VALID_JSON
+        + r  * WEIGHT_READ_RUNBOOK
+        + e  * WEIGHT_EFFICIENCY
+        for f, nw, v, r, e in zip(final, no_wrong, vj, rb, eff)
+    ]
+
+
+# Split mode: TRL sums five separate signals (preferred — gives per-signal trackio curves)
 ALL_REWARD_FUNCS = [
     reward_final_score,
     reward_no_wrong_channels,
@@ -87,3 +144,6 @@ ALL_REWARD_FUNCS = [
     reward_read_runbook,
     reward_efficiency,
 ]
+
+# Combined mode: single function — use when kwargs aren't flowing in split mode
+COMBINED_REWARD_FUNC = [combined_reward]

training/rollout.py

@@ -282,6 +282,20 @@ def rollout_once(
         # No posts at all — no spray happened (reward_final_score handles missing notif)
         no_wrong_channels = 1.0
 
+    read_runbook_reward = 1.0 if read_runbook_done else 0.0
+    combined = (
+        final_score       * 0.45
+        + no_wrong_channels * 0.15
+        + valid_json_ratio  * 0.15
+        + read_runbook_reward * 0.15
+        + efficiency        * 0.10
+    )
+    print(
+        f"[rollout] steps={step} final={final_score:.3f} "
+        f"json={valid_json_ratio:.2f} runbook={read_runbook_reward:.0f} "
+        f"no_wrong={no_wrong_channels:.2f} eff={efficiency:.2f} "
+        f"→ combined={combined:.3f}"
+    )
     return {
         "prompt_ids": prompt_ids,
         "completion_ids": completion_ids,
@@ -289,7 +303,7 @@ def rollout_once(
         "final_score_reward": final_score,
         "no_wrong_channels_reward": no_wrong_channels,
         "valid_json_reward": valid_json_ratio,
-        "read_runbook_reward": 1.0 if read_runbook_done else 0.0,
+        "read_runbook_reward": read_runbook_reward,
         "efficiency_reward": efficiency,
     }
 

training/train_v2.ipynb

@@ -0,0 +1,675 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# PM-Ops GRPO Training v2\n",
+    "\n",
+    "### Fixes vs v1\n",
+    "| Bug | Root cause | Fix |\n",
+    "|-----|-----------|-----|\n",
+    "| All rewards = 0 | TRL 1.3.0.dev0 (git) broke the `rollout_func` kwargs mechanism | Pin TRL to **1.2.0** stable; single `reward` key |\n",
+    "| Only 4 training steps | T4 (14 GB) + `gradient_accumulation=64` = ~4 steps | Auto-detect GPU, halve accum on T4 |\n",
+    "| Silent failure | No logging when rewards are 0 | Rollout prints every episode; reward func warns on zeros |\n",
+    "| `PatchFastRL` on wrong TRL | Patched TRL 0.17, then TRL was upgraded | Don't call `PatchFastRL` at all |\n",
+    "\n",
+    "**GPU recommendation:** A100 40 GB (Colab Pro) → 150 steps, ~60 min  \n",
+    "T4 works but only yields ~18 steps (still useful for smoke-testing the reward pipeline)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 0. Install Dependencies"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Step 1: Unsloth first (installs a TRL it likes, ~0.18-0.24)\n",
+    "!pip install -q unsloth vllm\n",
+    "\n",
+    "# Step 2: Force TRL to 1.2.0 stable (pip will warn about Unsloth conflict — safe to ignore;\n",
+    "#          Unsloth 2026.4.x patches work with TRL 1.0+ despite the metadata claim)\n",
+    "!pip install -q \"trl==1.2.0\" --force-reinstall\n",
+    "\n",
+    "# Step 3: Supporting packages\n",
+    "!pip install -q peft accelerate \"datasets>=3.4.1,<4.4.0\"\n",
+    "!pip install -q \"openenv-core>=0.2.2\" fastapi \"uvicorn[standard]\" \"pydantic>=2.0.0\"\n",
+    "!pip install -q trackio\n",
+    "\n",
+    "print('Done — restart kernel, then run all cells from top.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 1. Version Check + GPU Detect"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch, trl, trl.experimental.openenv\n",
+    "\n",
+    "print(f'TRL:    {trl.__version__}  (need 1.2.0)')\n",
+    "assert trl.__version__.startswith('1.2'), f'Wrong TRL version: {trl.__version__}. Rerun install cell and restart kernel.'\n",
+    "print('openenv: OK')\n",
+    "\n",
+    "gpu = torch.cuda.get_device_properties(0)\n",
+    "TOTAL_GB = round(gpu.total_memory / 1024**3, 1)\n",
+    "IS_A100  = TOTAL_GB >= 35\n",
+    "print(f'GPU:    {gpu.name}  ({TOTAL_GB} GB)')\n",
+    "print(f'Mode:   {\"A100 (full config)\" if IS_A100 else \"T4 (reduced config — smoke-test only)\"}')\n",
+    "\n",
+    "# Auto-adapted config values\n",
+    "GRAD_ACCUM   = 32  if IS_A100 else 8\n",
+    "NUM_GEN      = 4   if IS_A100 else 2\n",
+    "MAX_COMP_LEN = 512 if IS_A100 else 256\n",
+    "GPU_MEM_UTIL = 0.6 if IS_A100 else 0.5\n",
+    "USE_VLLM     = IS_A100  # vLLM colocate only safe on A100\n",
+    "\n",
+    "print(f'\\nConfig: grad_accum={GRAD_ACCUM}, num_gen={NUM_GEN}, max_comp={MAX_COMP_LEN}, vllm={USE_VLLM}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 2. Clone PM-Ops Repo"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os, sys\n",
+    "\n",
+    "REPO_URL = 'https://huggingface.co/spaces/TheCrustaceans/Pm-ops'\n",
+    "REPO_DIR = '/content/Pm_ops'\n",
+    "\n",
+    "if not os.path.exists(REPO_DIR):\n",
+    "    !git clone --depth=1 -q {REPO_URL} {REPO_DIR}\n",
+    "    print(f'Cloned → {REPO_DIR}')\n",
+    "else:\n",
+    "    # Pull latest — picks up the fixed rewards.py and rollout.py\n",
+    "    !git -C {REPO_DIR} pull -q\n",
+    "    print(f'Updated: {REPO_DIR}')\n",
+    "\n",
+    "for p in [REPO_DIR, os.path.join(REPO_DIR, 'training')]:\n",
+    "    if p not in sys.path:\n",
+    "        sys.path.insert(0, p)\n",
+    "\n",
+    "os.chdir(REPO_DIR)\n",
+    "print(f'CWD: {os.getcwd()}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 3. HuggingFace Login"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import notebook_login\n",
+    "notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 4. Start Local PM-Ops Server\n",
+    "Localhost removes ~200 ms/step network round-trip."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import subprocess, time, requests\n",
+    "\n",
+    "server_proc = subprocess.Popen(\n",
+    "    [sys.executable, '-m', 'uvicorn', 'server.app:app', '--host', '0.0.0.0', '--port', '8000'],\n",
+    "    cwd=REPO_DIR, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,\n",
+    ")\n",
+    "ENV_URL = 'http://localhost:8000'\n",
+    "\n",
+    "for i in range(30):\n",
+    "    try:\n",
+    "        if requests.get(f'{ENV_URL}/', timeout=2).status_code == 200:\n",
+    "            print(f'PM-Ops server ready (pid={server_proc.pid})')\n",
+    "            break\n",
+    "    except Exception:\n",
+    "        pass\n",
+    "    time.sleep(1)\n",
+    "else:\n",
+    "    raise RuntimeError('Server did not start in 30 s')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 5. Verify Environment"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openenv.core import GenericEnvClient\n",
+    "\n",
+    "with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
+    "    res = _env.reset()\n",
+    "    obs = res.observation if hasattr(res, 'observation') else res\n",
+    "    brief = getattr(obs, 'task_brief', '') or (obs.get('task_brief', '') if isinstance(obs, dict) else '')\n",
+    "    print(f'task_brief: {brief[:100]}...')\n",
+    "    _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
+    "    print('Env step OK')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 6. Load Model with Unsloth\n",
+    "\n",
+    "**No `PatchFastRL`** — that call patches TRL internals and is incompatible with TRL 1.2.0.  \n",
+    "Unsloth is still used for 4-bit loading + LoRA, which work independently of the TRL version."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from unsloth import FastLanguageModel\n",
+    "\n",
+    "MODEL_NAME  = 'unsloth/Qwen3-1.7B-unsloth-bnb-4bit'  # pre-quantised, faster download\n",
+    "MAX_SEQ_LEN = 4096 + MAX_COMP_LEN\n",
+    "LORA_RANK   = 16\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name=MODEL_NAME,\n",
+    "    max_seq_length=MAX_SEQ_LEN,\n",
+    "    load_in_4bit=True,\n",
+    "    fast_inference=False,          # keep False — vLLM colocate only on A100 via GRPOConfig\n",
+    "    max_lora_rank=LORA_RANK,\n",
+    "    gpu_memory_utilization=GPU_MEM_UTIL,\n",
+    ")\n",
+    "\n",
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r=LORA_RANK,\n",
+    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
+    "    lora_alpha=LORA_RANK,\n",
+    "    use_gradient_checkpointing='unsloth',\n",
+    "    random_state=42,\n",
+    ")\n",
+    "tokenizer.pad_token = tokenizer.eos_token\n",
+    "print(f'Model ready: {MODEL_NAME}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 7. Generate Training Dataset"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import Dataset\n",
+    "from training.dataset import generate_triage_dataset\n",
+    "\n",
+    "N_EPISODES = 150\n",
+    "rows = generate_triage_dataset(n_episodes=N_EPISODES, base_seed=42)\n",
+    "# 'prompt' column = seed-embedded string; GRPOTrainer passes it as-is to rollout_func\n",
+    "dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
+    "print(f'Dataset: {len(dataset)} episodes')\n",
+    "print(f'Sample:  {dataset[0][\"prompt\"][:100]}...')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 8. Rollout Function + Reward\n",
+    "\n",
+    "**Key change from v1:** Instead of returning 5 separate reward keys (which broke in TRL 1.3.0.dev0),\n",
+    "we return a single `\"reward\"` key with the pre-weighted combined scalar.  \n",
+    "The reward function is a trivial passthrough — no kwargs plumbing can fail."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from openenv.core import GenericEnvClient\n",
+    "from training.rollout import rollout_once\n",
+    "\n",
+    "# ── Persistent env connection for training ────────────────────────────────────\n",
+    "sync_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
+    "sync_env.connect()\n",
+    "print('Training env connected')\n",
+    "\n",
+    "TRAIN_MAX_STEPS = 15\n",
+    "\n",
+    "# ── Combined rollout (single reward key) ─────────────────────────────────────\n",
+    "def make_rollout_func(env, tok, max_steps=TRAIN_MAX_STEPS):\n",
+    "    \"\"\"\n",
+    "    Returns a rollout_func compatible with TRL 1.2.0 GRPOTrainer.\n",
+    "\n",
+    "    Returns a single 'reward' key (weighted sum of all signals) so the reward\n",
+    "    function is a trivial passthrough — avoids the kwargs-flow bug from v1.\n",
+    "    \"\"\"\n",
+    "    def rollout_func(prompts, trainer=None):\n",
+    "        out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n",
+    "        for prompt in prompts:\n",
+    "            ep = rollout_once(\n",
+    "                trainer=trainer, sync_env=env, tokenizer=tok,\n",
+    "                dataset_prompt=prompt, max_steps=max_steps,\n",
+    "            )\n",
+    "            # Weighted combination (same weights as rewards.py)\n",
+    "            combined = (\n",
+    "                ep['final_score_reward']       * 0.45 +\n",
+    "                ep['no_wrong_channels_reward'] * 0.15 +\n",
+    "                ep['valid_json_reward']        * 0.15 +\n",
+    "                ep['read_runbook_reward']      * 0.15 +\n",
+    "                ep['efficiency_reward']        * 0.10\n",
+    "            )\n",
+    "            out['prompt_ids'].append(ep['prompt_ids'])\n",
+    "            out['completion_ids'].append(ep['completion_ids'])\n",
+    "            out['logprobs'].append(ep['logprobs'])\n",
+    "            out['reward'].append(combined)\n",
+    "        return out\n",
+    "    return rollout_func\n",
+    "\n",
+    "rollout_func = make_rollout_func(sync_env, tokenizer)\n",
+    "\n",
+    "# ── Reward function: trivial passthrough ─────────────────────────────────────\n",
+    "def reward_func(completions, **kwargs):\n",
+    "    \"\"\"\n",
+    "    Reads pre-computed combined reward from rollout kwargs.\n",
+    "    Warns loudly if kwargs are empty (kwargs-flow is broken).\n",
+    "    \"\"\"\n",
+    "    n = len(completions)\n",
+    "    rewards = kwargs.get('reward', [])\n",
+    "    if not rewards:\n",
+    "        print(\n",
+    "            f'[ERROR] reward_func received empty kwargs! '\n",
+    "            f'kwargs keys: {list(kwargs.keys())}. '\n",
+    "            f'All rewards will be 0 — training is broken. '\n",
+    "            f'TRL version: {trl.__version__}'\n",
+    "        )\n",
+    "        return [0.0] * n\n",
+    "    return [float(r) for r in rewards]\n",
+    "\n",
+    "print('rollout_func and reward_func defined')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 9. Smoke Test — Verify Reward Signal Before Training\n",
+    "\n",
+    "Run one episode and confirm:\n",
+    "1. Rollout completes without error\n",
+    "2. `reward` key exists in rollout output\n",
+    "3. `read_runbook_reward` = 1.0 (step_aware_fallback always calls read_runbook at step 0)\n",
+    "4. `reward_func` receives non-empty kwargs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print('── Smoke test ──────────────────────────────')\n",
+    "\n",
+    "test_prompt = dataset[0]['prompt']\n",
+    "print(f'Prompt: {test_prompt[:80]}...')\n",
+    "\n",
+    "test_out = rollout_func([test_prompt], trainer=None)\n",
+    "\n",
+    "assert 'reward' in test_out,            'FAIL: rollout missing reward key'\n",
+    "assert len(test_out['reward']) == 1,    'FAIL: reward list wrong length'\n",
+    "assert len(test_out['prompt_ids']) == 1, 'FAIL: prompt_ids wrong length'\n",
+    "\n",
+    "r = test_out['reward'][0]\n",
+    "n_comp_tokens = len(test_out['completion_ids'][0])\n",
+    "print(f'combined reward  : {r:.4f}')\n",
+    "print(f'completion tokens: {n_comp_tokens}')\n",
+    "\n",
+    "# Verify reward_func receives the kwarg\n",
+    "fake_completions = ['fake completion']\n",
+    "rf_out = reward_func(fake_completions, reward=[r])\n",
+    "assert rf_out == [r], 'FAIL: reward_func passthrough broken'\n",
+    "\n",
+    "if r == 0.0:\n",
+    "    print('\\n[WARN] combined reward = 0.0 on first episode.')\n",
+    "    print('  This is OK if it is the very first step (model learning).')\n",
+    "    print('  Check rollout log above for individual signal breakdown.')\n",
+    "else:\n",
+    "    print(f'\\nSmoke test PASSED — reward = {r:.4f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 10. Configure GRPO Training"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import GRPOConfig\n",
+    "\n",
+    "OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v2'\n",
+    "HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n",
+    "\n",
+    "grpo_config = GRPOConfig(\n",
+    "    # ── Training ──────────────────────────────────────────────────────────────\n",
+    "    num_train_epochs=1,\n",
+    "    learning_rate=5e-6,\n",
+    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
+    "    per_device_train_batch_size=1,\n",
+    "    warmup_steps=5,\n",
+    "    num_generations=NUM_GEN,\n",
+    "\n",
+    "    # ── Sequence lengths ──────────────────────────────────────────────────────\n",
+    "    max_completion_length=MAX_COMP_LEN,\n",
+    "    max_prompt_length=4096,\n",
+    "\n",
+    "    # ── Generation ────────────────────────────────────────────────────────────\n",
+    "    # use_vllm=True only on A100 (requires fast_inference=True in FastLanguageModel)\n",
+    "    # On T4: use_vllm=False → standard HF generate, stop tokens respected via eos_token\n",
+    "    use_vllm=USE_VLLM,\n",
+    "\n",
+    "    # ── Output + logging ──────────────────────────────────────────────────────\n",
+    "    output_dir=OUTPUT_DIR,\n",
+    "    report_to='trackio',\n",
+    "    trackio_space_id=OUTPUT_DIR,\n",
+    "    logging_steps=1,\n",
+    "    save_steps=25,\n",
+    "    gradient_checkpointing=False,  # Unsloth handles this via get_peft_model\n",
+    ")\n",
+    "\n",
+    "total_steps = len(dataset) // (grpo_config.per_device_train_batch_size * GRAD_ACCUM)\n",
+    "print(f'Output dir : {OUTPUT_DIR}')\n",
+    "print(f'Total steps: ~{total_steps}')\n",
+    "print(f'Eff. batch : {grpo_config.per_device_train_batch_size * GRAD_ACCUM}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 11. Create Trainer"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from trl import GRPOTrainer\n",
+    "\n",
+    "trainer = GRPOTrainer(\n",
+    "    model=model,\n",
+    "    processing_class=tokenizer,\n",
+    "    reward_funcs=reward_func,       # single function, not a list\n",
+    "    train_dataset=dataset,\n",
+    "    args=grpo_config,\n",
+    "    rollout_func=rollout_func,\n",
+    ")\n",
+    "print('GRPOTrainer ready')\n",
+    "print(f'vLLM engine present: {hasattr(model, \"vllm_engine\")}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 12. Train\n",
+    "\n",
+    "**Watch the Colab output for rollout logs:**\n",
+    "```\n",
+    "[rollout] steps=5 final=0.720 json=0.80 runbook=1 no_wrong=1.00 eff=0.67 → combined=0.674\n",
+    "```\n",
+    "If `combined > 0` appears in rollout logs but `train/reward = 0.0` in trackio → TRL still not passing kwargs.  \n",
+    "If `combined > 0` appears in both → training is working.\n",
+    "\n",
+    "**Key metrics to watch in trackio:**\n",
+    "- `train/reward` — must be > 0, should rise over training\n",
+    "- `train/loss` — must be non-zero (proves GRPO advantage ≠ 0)\n",
+    "- `train/grad_norm` — must be non-zero (proves weight updates are happening)\n",
+    "- `train/frac_reward_zero_std` — should drop below 1.0 quickly"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer_stats = trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Training summary\n",
+    "used_gb = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
+    "train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n",
+    "print(f'Training time : {train_mins} min')\n",
+    "print(f'Peak GPU usage: {used_gb} GB / {TOTAL_GB} GB ({round(used_gb/TOTAL_GB*100,1)}%)')\n",
+    "print(f'Final loss    : {trainer_stats.metrics.get(\"train_loss\", \"n/a\")}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 13. Save + Push to HuggingFace"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sync_env.close()  # close training connection before saving\n",
+    "\n",
+    "# Merge LoRA into bf16 — safe Unsloth path (never call trainer.save_model() directly)\n",
+    "model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\n",
+    "print(f'Saved → {OUTPUT_DIR}')\n",
+    "\n",
+    "model.push_to_hub_merged(HF_REPO_ID, tokenizer, save_method='merged_16bit')\n",
+    "print(f'Pushed → https://huggingface.co/{HF_REPO_ID}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 14. Evaluate: Baseline vs Trained\n",
+    "\n",
+    "Uses the **already-loaded model** (via `FastLanguageModel.for_inference`) — no reload needed.  \n",
+    "Connects to the remote HF Space for a clean eval (not localhost)."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from unsloth import FastLanguageModel\n",
+    "from training.rollout import extract_json_action, step_aware_fallback, build_messages, _obs_to_dict\n",
+    "from training.rollout import _current_obs_text\n",
+    "from inference import baseline_agent\n",
+    "\n",
+    "EVAL_URL      = 'https://adityaguntur-pm-ops.hf.space'\n",
+    "N_EVAL        = 10\n",
+    "EVAL_MAX_STEPS = 15\n",
+    "\n",
+    "FastLanguageModel.for_inference(model)  # switch to fast inference mode\n",
+    "\n",
+    "\n",
+    "def eval_trained(env, mdl, tok, n=N_EVAL):\n",
+    "    scores = []\n",
+    "    for i in range(n):\n",
+    "        result = env.reset()\n",
+    "        obs    = result.observation if hasattr(result, 'observation') else result\n",
+    "        obs_dict   = _obs_to_dict(obs)\n",
+    "        task_brief = obs_dict.get('task_brief', '')\n",
+    "        history, step, score, done = [], 0, 0.0, False\n",
+    "\n",
+    "        while not done and step < EVAL_MAX_STEPS:\n",
+    "            obs_text  = _current_obs_text(obs_dict, step, task_brief)\n",
+    "            msgs      = build_messages(history, obs_text)\n",
+    "            prompt_t  = tok.apply_chat_template(\n",
+    "                msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
+    "            )\n",
+    "            inputs    = tok([prompt_t], return_tensors='pt').to(mdl.device)\n",
+    "            with torch.no_grad():\n",
+    "                out_ids = mdl.generate(**inputs, max_new_tokens=512, temperature=0.0)\n",
+    "            completion = tok.decode(out_ids[0][len(inputs.input_ids[0]):], skip_special_tokens=True)\n",
+    "\n",
+    "            parsed  = extract_json_action(completion) or step_aware_fallback(step)\n",
+    "            result  = env.step({'action_type': parsed['action_type'], 'args': parsed.get('args', {})})\n",
+    "            obs_d2  = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "            obs_dict = obs_d2\n",
+    "            done    = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
+    "            score   = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
+    "            history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n",
+    "            step += 1\n",
+    "\n",
+    "        scores.append(score)\n",
+    "        print(f'  Trained  ep {i+1}/{n}: score={score:.3f}')\n",
+    "    return scores\n",
+    "\n",
+    "\n",
+    "def eval_baseline(env, n=N_EVAL):\n",
+    "    scores = []\n",
+    "    for i in range(n):\n",
+    "        result   = env.reset()\n",
+    "        obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "        org_cfg, done, step, score = {}, False, 0, 0.0\n",
+    "        while not done and step < EVAL_MAX_STEPS:\n",
+    "            at, args = baseline_agent(obs_dict, org_cfg)\n",
+    "            result   = env.step({'action_type': at, 'args': args})\n",
+    "            obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "            done  = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
+    "            score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
+    "            step += 1\n",
+    "        scores.append(score)\n",
+    "        print(f'  Baseline ep {i+1}/{n}: score={score:.3f}')\n",
+    "    return scores\n",
+    "\n",
+    "\n",
+    "print('─── Baseline ───')\n",
+    "with GenericEnvClient(base_url=EVAL_URL).sync() as env_eval:\n",
+    "    baseline_scores = eval_baseline(env_eval)\n",
+    "\n",
+    "print('─── Trained ────')\n",
+    "with GenericEnvClient(base_url=EVAL_URL).sync() as env_eval:\n",
+    "    trained_scores = eval_trained(env_eval, model, tokenizer)\n",
+    "\n",
+    "print(f'\\nBaseline avg : {sum(baseline_scores)/N_EVAL:.3f}')\n",
+    "print(f'Trained  avg : {sum(trained_scores)/N_EVAL:.3f}')\n",
+    "print(f'Improvement  : +{(sum(trained_scores)-sum(baseline_scores))/N_EVAL:.3f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 15. Plot Results"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
+    "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
+    "\n",
+    "ax, x, w = axes[0], np.arange(N_EVAL), 0.35\n",
+    "ax.bar(x - w/2, baseline_scores, w, label='Baseline', color='steelblue', alpha=0.8)\n",
+    "ax.bar(x + w/2, trained_scores,  w, label='GRPO v2',  color='coral',     alpha=0.8)\n",
+    "ax.axhline(sum(baseline_scores)/N_EVAL, color='steelblue', linestyle='--', alpha=0.5, linewidth=1)\n",
+    "ax.axhline(sum(trained_scores)/N_EVAL,  color='coral',     linestyle='--', alpha=0.5, linewidth=1)\n",
+    "ax.set(xlabel='Episode', ylabel='Reward', title='Per-episode: Baseline vs GRPO v2',\n",
+    "       xticks=x, ylim=(0, 1.05))\n",
+    "ax.legend()\n",
+    "\n",
+    "ax2 = axes[1]\n",
+    "avgs = [sum(baseline_scores)/N_EVAL, sum(trained_scores)/N_EVAL]\n",
+    "bars = ax2.bar(['Baseline', 'GRPO v2'], avgs, color=['steelblue', 'coral'], alpha=0.85, width=0.5)\n",
+    "for bar, val in zip(bars, avgs):\n",
+    "    ax2.text(bar.get_x() + bar.get_width()/2, val + 0.01, f'{val:.3f}',\n",
+    "             ha='center', fontsize=13, fontweight='bold')\n",
+    "ax2.set(ylabel='Average reward', title=f'Average over {N_EVAL} episodes', ylim=(0, 1.05))\n",
+    "\n",
+    "plt.tight_layout()\n",
+    "plt.savefig('eval_results_v2.png', dpi=150, bbox_inches='tight')\n",
+    "plt.show()\n",
+    "print('Saved: eval_results_v2.png')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": ["## 16. Teardown"]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "server_proc.terminate()\n",
+    "print('Local PM-Ops server stopped')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.12.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}