feat(training): add clean v4 notebook with all reward fixes

Single-pass, no dead code. Key improvements over v3/pink notebooks:
- gen_slot correctly offsets env seed per GRPO generation
- runbook-compliance reward (label/priority/team/channel vs org config)
- compute_rollout_reward imported from training/rewards.py
- temperature 1.1 + top_k=50 for diverse rollouts
- PatchFastRL called once before trl imports
- No duplicate git-pull or dead rollout cells

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

training/train_v4.ipynb

@@ -0,0 +1,909 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
+  "language_info": {"name": "python", "version": "3.12.0"}
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "v4-title",
+   "metadata": {},
+   "source": [
+    "# PM-Ops GRPO Training v4\n",
+    "\n",
+    "Clean rewrite — all known v3 bugs fixed.\n",
+    "\n",
+    "| Fix | Details |\n",
+    "|---|---|\n",
+    "| Constant reward=0.350 | Replaced env_score formula with runbook-compliance scoring |\n",
+    "| `gen_slot=0` override bug | gen_slot now correctly offsets env seed per GRPO generation |\n",
+    "| Double `env.step` per step | Removed duplicate step call inside try/except |\n",
+    "| Dataset/env task-type mismatch | `dataset.py` pre-filters seeds to triage-only episodes |\n",
+    "| Near-greedy generation | Temperature 1.1 + top_k=50 for diverse rollouts |\n",
+    "| PatchFastRL called twice | Called once, unconditionally, before any trl imports |\n",
+    "\n",
+    "**Reward**: read_runbook (+0.10) + valid_label (±0.20/0.10) + valid_priority (±0.15/0.10) + valid_team (±0.20/0.10) + right_channel (±0.25/0.10) + env_bonus (×0.10). Varies per org config — different valid values per seed.  \n",
+    "**Stack**: Unsloth 2026.x · TRL 0.22.2 · Qwen3-1.7B · PMOpsGRPOTrainer  \n",
+    "**GPU**: A100-40GB → ~2 hrs (SFT 15 min + GRPO 90 min)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s0",
+   "metadata": {},
+   "source": ["## 0. Install"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-install",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "%%capture\n",
+    "import os\n",
+    "!pip install --upgrade -qqq uv\n",
+    "if 'COLAB_' not in ''.join(os.environ.keys()):\n",
+    "    !uv pip install unsloth vllm\n",
+    "else:\n",
+    "    import subprocess\n",
+    "    is_t4 = 'Tesla T4' in str(subprocess.check_output(['nvidia-smi']))\n",
+    "    _vllm   = 'vllm==0.9.2'   if is_t4 else 'vllm==0.15.1'\n",
+    "    _triton = 'triton==3.2.0' if is_t4 else 'triton'\n",
+    "    !uv pip install -qqq --upgrade {_vllm} torchvision bitsandbytes xformers unsloth\n",
+    "    !uv pip install -qqq {_triton}\n",
+    "!uv pip install transformers==4.56.2\n",
+    "!uv pip install --no-deps trl==0.22.2\n",
+    "!pip install 'numpy==1.26.4' --break-system-packages -q\n",
+    "!pip install openenv openenv-core -q"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s1",
+   "metadata": {},
+   "source": ["## 1. GPU Config + Patch"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-config",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "from unsloth import FastLanguageModel, PatchFastRL\n",
+    "\n",
+    "# Must patch BEFORE any trl imports\n",
+    "PatchFastRL('GRPO', FastLanguageModel)\n",
+    "\n",
+    "import trl\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",
+    "\n",
+    "NUM_GEN         = 6   if IS_A100 else 2\n",
+    "GRAD_ACCUM      = 32  if IS_A100 else 8\n",
+    "MAX_COMP_LEN    = 384\n",
+    "N_SFT_EPISODES  = 120 if IS_A100 else 20\n",
+    "N_GRPO_EPISODES = 150 if IS_A100 else 30\n",
+    "TRAIN_MAX_STEPS = 12\n",
+    "\n",
+    "print(f'torch={torch.__version__}  trl={trl.__version__}')\n",
+    "print(f'GPU: {gpu.name} ({TOTAL_GB} GB)  IS_A100={IS_A100}')\n",
+    "print(f'num_gen={NUM_GEN}  grad_accum={GRAD_ACCUM}  '\n",
+    "      f'sft_eps={N_SFT_EPISODES}  grpo_eps={N_GRPO_EPISODES}  max_steps={TRAIN_MAX_STEPS}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s2",
+   "metadata": {},
+   "source": ["## 2. Clone PM-Ops Repo"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-clone",
+   "metadata": {},
+   "execution_count": null,
+   "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",
+    "    !git -C {REPO_DIR} pull -q origin main\n",
+    "    print(f'Pulled -> {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",
+    "os.chdir(REPO_DIR)\n",
+    "print(f'CWD: {os.getcwd()}')\n",
+    "!git -C {REPO_DIR} log --oneline -3"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s3",
+   "metadata": {},
+   "source": ["## 3. HuggingFace Login"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-hf-login",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from huggingface_hub import notebook_login\n",
+    "notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s4",
+   "metadata": {},
+   "source": ["## 4. Start PM-Ops Server"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-server",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import subprocess, time, requests\n",
+    "\n",
+    "# Kill any leftover server from a previous run\n",
+    "subprocess.run(['pkill', '-f', 'uvicorn'], capture_output=True)\n",
+    "time.sleep(1)\n",
+    "\n",
+    "server_proc = subprocess.Popen(\n",
+    "    [sys.executable, '-m', 'uvicorn', 'server.app:app',\n",
+    "     '--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 _ 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",
+   "id": "v4-s5",
+   "metadata": {},
+   "source": ["## 5. Verify Env"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-verify-env",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from openenv.core import GenericEnvClient\n",
+    "from training.rollout import _obs_to_dict\n",
+    "\n",
+    "with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
+    "    r   = _env.reset(seed=42)\n",
+    "    obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
+    "    print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n",
+    "    step_r = _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
+    "    rb_obs = _obs_to_dict(step_r.observation if hasattr(step_r, 'observation') else step_r)\n",
+    "    data   = (rb_obs.get('last_action_result') or {}).get('data', {})\n",
+    "    org    = data.get('org_config', {}) if isinstance(data, dict) else {}\n",
+    "    print(f'org labels    : {list(org.get(\"label_taxonomy\", {}).values())}')\n",
+    "    print(f'org priorities: {org.get(\"priority_levels\", [])}')\n",
+    "    print(f'org channels  : {list(org.get(\"oncall_channels\", {}).values())}')\n",
+    "    print('Env verify: OK')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s6",
+   "metadata": {},
+   "source": ["## 6. Load Model — Unsloth 4-bit + LoRA"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-model",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
+    "LORA_RANK  = 16\n",
+    "\n",
+    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+    "    model_name             = MODEL_NAME,\n",
+    "    max_seq_length         = 4096 + MAX_COMP_LEN,\n",
+    "    load_in_4bit           = True,\n",
+    "    fast_inference         = False,\n",
+    "    max_lora_rank          = LORA_RANK,\n",
+    "    gpu_memory_utilization = 0.55,\n",
+    ")\n",
+    "model = FastLanguageModel.get_peft_model(\n",
+    "    model,\n",
+    "    r                          = LORA_RANK,\n",
+    "    target_modules             = ['q_proj','k_proj','v_proj','o_proj',\n",
+    "                                   '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",
+    "tokenizer.padding_side = 'left'\n",
+    "model.print_trainable_parameters()\n",
+    "print(f'GPU after load: {round(torch.cuda.max_memory_reserved()/1024**3, 2)} GB / {TOTAL_GB} GB')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-phase1",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## Phase 1 — SFT Warmup\n",
+    "\n",
+    "Teach the model the JSON output format and PM-ops workflow before GRPO.  \n",
+    "`baseline_agent` runs N deterministic episodes → ~6 steps each → supervised (prompt, completion) pairs.  \n",
+    "2 SFT epochs (~15 min on A100). Without SFT, all rollouts output freeform text → zero GRPO gradient."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s7",
+   "metadata": {},
+   "source": ["## 7. Generate SFT Dataset"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-sft-data",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import json as _json\n",
+    "from datasets import Dataset\n",
+    "from inference import baseline_agent\n",
+    "from training.rollout import _obs_to_dict, _current_obs_text, build_messages\n",
+    "\n",
+    "\n",
+    "def generate_sft_dataset(env_url, tok, n_episodes, seed_start=2000):\n",
+    "    examples = []\n",
+    "    with GenericEnvClient(base_url=env_url).sync() as env:\n",
+    "        for i in range(n_episodes):\n",
+    "            result     = env.reset(seed=seed_start + i)\n",
+    "            obs_dict   = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "            task_brief = obs_dict.get('task_brief', '')\n",
+    "            turn_history, org_config, step, done = [], {}, 0, False\n",
+    "\n",
+    "            while not done and step < 8:\n",
+    "                obs_text          = _current_obs_text(obs_dict, step, task_brief)\n",
+    "                action_type, args = baseline_agent(obs_dict, org_config)\n",
+    "                payload           = {'action_type': action_type, 'args': args}\n",
+    "                completion        = '```json\\n' + _json.dumps(payload) + '\\n```'\n",
+    "                msgs              = build_messages(turn_history, obs_text)\n",
+    "                prompt            = tok.apply_chat_template(\n",
+    "                    msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
+    "                )\n",
+    "                examples.append({'text': prompt + completion + tok.eos_token})\n",
+    "                turn_history.append({\n",
+    "                    'obs_text': obs_text, 'completion': completion,\n",
+    "                    'is_runbook': (action_type == 'meta.read_runbook'),\n",
+    "                })\n",
+    "                result   = env.step({'action_type': action_type, 'args': args})\n",
+    "                obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "                if action_type == 'meta.read_runbook':\n",
+    "                    last = obs_dict.get('last_action_result') or {}\n",
+    "                    if last.get('ok'):\n",
+    "                        data = last.get('data') or {}\n",
+    "                        if isinstance(data, dict) and 'org_config' in data:\n",
+    "                            org_config.update(data['org_config'])\n",
+    "                done = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
+    "                step += 1\n",
+    "\n",
+    "            if (i + 1) % 20 == 0:\n",
+    "                print(f'  {i+1}/{n_episodes} eps — {len(examples)} examples')\n",
+    "\n",
+    "    return examples\n",
+    "\n",
+    "\n",
+    "print(f'Generating {N_SFT_EPISODES} SFT demonstration episodes...')\n",
+    "sft_raw     = generate_sft_dataset(ENV_URL, tokenizer, n_episodes=N_SFT_EPISODES)\n",
+    "sft_dataset = Dataset.from_list(sft_raw)\n",
+    "print(f'SFT dataset: {len(sft_dataset)} examples (~{len(sft_dataset)//6} eps × 6 steps)')\n",
+    "print(f'Sample (first 300 chars):\\n{sft_raw[0][\"text\"][:300]}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s8",
+   "metadata": {},
+   "source": ["## 8. SFT Training"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-sft-train",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from trl import SFTTrainer, SFTConfig\n",
+    "\n",
+    "sft_cfg = SFTConfig(\n",
+    "    dataset_text_field          = 'text',\n",
+    "    max_seq_length              = 2048,\n",
+    "    num_train_epochs            = 2,\n",
+    "    per_device_train_batch_size = 4,\n",
+    "    gradient_accumulation_steps = 4,\n",
+    "    learning_rate               = 2e-4,\n",
+    "    warmup_steps                = 10,\n",
+    "    output_dir                  = 'pm-ops-sft-warmup',\n",
+    "    report_to                   = 'none',\n",
+    "    logging_steps               = 5,\n",
+    "    save_strategy               = 'no',\n",
+    "    dataloader_num_workers      = 0,\n",
+    ")\n",
+    "sft_steps = (\n",
+    "    len(sft_dataset)\n",
+    "    // (sft_cfg.per_device_train_batch_size * sft_cfg.gradient_accumulation_steps)\n",
+    "    * sft_cfg.num_train_epochs\n",
+    ")\n",
+    "print(f'SFT: {len(sft_dataset)} examples × {sft_cfg.num_train_epochs} epochs → ~{sft_steps} steps')\n",
+    "\n",
+    "sft_trainer = SFTTrainer(model=model, tokenizer=tokenizer,\n",
+    "                         train_dataset=sft_dataset, args=sft_cfg)\n",
+    "sft_stats = sft_trainer.train()\n",
+    "loss    = sft_stats.metrics.get('train_loss', 0)\n",
+    "runtime = sft_stats.metrics.get('train_runtime', 0)\n",
+    "print(f'SFT done: {round(runtime/60, 1)} min  loss={loss:.3f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s9",
+   "metadata": {},
+   "source": ["## 9. Verify SFT — Model Must Output Valid JSON"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-sft-verify",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from training.rollout import extract_json_action\n",
+    "\n",
+    "model.eval()\n",
+    "with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
+    "    r        = _env.reset(seed=99001)\n",
+    "    obs_dict = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
+    "    obs_text = _current_obs_text(obs_dict, 0, obs_dict.get('task_brief', ''))\n",
+    "    msgs     = build_messages([], obs_text)\n",
+    "    prompt   = tokenizer.apply_chat_template(\n",
+    "        msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
+    "    )\n",
+    "\n",
+    "inputs = tokenizer([prompt], return_tensors='pt').to(model.device)\n",
+    "with torch.no_grad():\n",
+    "    out = model.generate(**inputs, max_new_tokens=128, do_sample=False,\n",
+    "                         pad_token_id=tokenizer.eos_token_id)\n",
+    "completion = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
+    "parsed     = extract_json_action(completion)\n",
+    "\n",
+    "print(f'Output: {completion[:400]}')\n",
+    "print(f'Parsed: {parsed}')\n",
+    "if parsed is not None:\n",
+    "    print('PASS ✓ — model outputs valid JSON after SFT')\n",
+    "else:\n",
+    "    print('FAIL ✗ — still no valid JSON. Run SFT again with more episodes or epochs.')\n",
+    "model.train()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-phase2",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## Phase 2 — GRPO\n",
+    "\n",
+    "Reward components — sum ≈ 0.90 when all correct:\n",
+    "\n",
+    "| Signal | Weight | Description |\n",
+    "|---|---|---|\n",
+    "| `read_runbook` | +0.10 | Did agent read the runbook first? |\n",
+    "| `valid_label` | +0.20 / −0.10 | Ticket label in org's `label_taxonomy`? |\n",
+    "| `valid_priority` | +0.15 / −0.10 | Ticket priority in org's `priority_levels`? |\n",
+    "| `valid_team` | +0.20 / −0.10 | Assigned team in org's `team_map`? |\n",
+    "| `right_channel` | +0.25 / −0.10/ch | Posted to org's `oncall_channels`? |\n",
+    "| `env_bonus` | ×0.10 | Env grader confirmation |\n",
+    "\n",
+    "Reward **varies per org config** (different valid values per seed) → genuine GRPO advantage signal.  \n",
+    "Each of the N GRPO generations gets a different env seed via `gen_slot` offset."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s10",
+   "metadata": {},
+   "source": ["## 10. GRPO Training Dataset"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-grpo-data",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from training.dataset import generate_triage_dataset\n",
+    "\n",
+    "rows         = generate_triage_dataset(n_episodes=N_GRPO_EPISODES, base_seed=42)\n",
+    "grpo_dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
+    "print(f'GRPO dataset: {len(grpo_dataset)} triage episodes')\n",
+    "print(f'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\n",
+    "print(f'Sample: {rows[0][\"prompt\"][:120]}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s11",
+   "metadata": {},
+   "source": ["## 11. GRPO Rollout + Reward"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-rollout",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import torch.nn.functional as F\n",
+    "from training.rollout import (\n",
+    "    _obs_to_dict, _current_obs_text, build_messages,\n",
+    "    extract_json_action, step_aware_fallback,\n",
+    ")\n",
+    "from training.dataset import parse_seed_from_prompt\n",
+    "from training.rewards import compute_rollout_reward\n",
+    "\n",
+    "grpo_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
+    "grpo_env.connect()\n",
+    "print('GRPO env connected')\n",
+    "\n",
+    "\n",
+    "def run_grpo_episode(trainer, env, tok, dataset_prompt, max_steps=TRAIN_MAX_STEPS, gen_slot=0):\n",
+    "    \"\"\"One full PM-ops episode. gen_slot offsets seed so parallel GRPO generations\n",
+    "    explore distinct env episodes even for the same base prompt.\"\"\"\n",
+    "    seed   = parse_seed_from_prompt(dataset_prompt)\n",
+    "    result = env.reset(seed=seed + gen_slot) if seed is not None else env.reset()\n",
+    "\n",
+    "    obs_dict   = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "    task_brief = obs_dict.get('task_brief') or dataset_prompt\n",
+    "\n",
+    "    prompt_ids, completion_ids, logprobs = [], [], []\n",
+    "    turn_history     = []\n",
+    "    valid_json_count = 0\n",
+    "    env_score        = 0.0\n",
+    "    step, done       = 0, False\n",
+    "\n",
+    "    # Runbook-compliance tracking\n",
+    "    read_runbook_done = False\n",
+    "    valid_labels:     set      = set()\n",
+    "    valid_priorities: set      = set()\n",
+    "    valid_teams:      set      = set()\n",
+    "    oncall_channels:  set      = set()\n",
+    "    ticket_label:    str | None = None\n",
+    "    ticket_priority: str | None = None\n",
+    "    assigned_team:   str | None = None\n",
+    "    posted_channels: list       = []\n",
+    "\n",
+    "    _model  = (trainer.accelerator.unwrap_model(trainer.model)\n",
+    "               if hasattr(trainer, 'accelerator') else trainer.model)\n",
+    "    _device = (trainer.accelerator.device\n",
+    "               if hasattr(trainer, 'accelerator') else next(_model.parameters()).device)\n",
+    "\n",
+    "    while not done and step < max_steps:\n",
+    "        obs_text    = _current_obs_text(obs_dict, step, task_brief)\n",
+    "        msgs        = build_messages(turn_history, obs_text)\n",
+    "        prompt_text = tok.apply_chat_template(\n",
+    "            msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
+    "        )\n",
+    "\n",
+    "        enc  = tok(prompt_text, return_tensors='pt', truncation=True, max_length=4096).to(_device)\n",
+    "        plen = enc['input_ids'].shape[1]\n",
+    "        with torch.no_grad():\n",
+    "            out = _model.generate(\n",
+    "                **enc,\n",
+    "                max_new_tokens         = MAX_COMP_LEN,\n",
+    "                do_sample              = True,\n",
+    "                temperature            = 1.1,\n",
+    "                top_p                  = 0.95,\n",
+    "                top_k                  = 50,\n",
+    "                pad_token_id           = tok.pad_token_id or tok.eos_token_id,\n",
+    "                output_scores          = True,\n",
+    "                return_dict_in_generate = True,\n",
+    "            )\n",
+    "\n",
+    "        cids            = out.sequences[0][plen:].tolist()\n",
+    "        completion_text = tok.decode(cids, skip_special_tokens=True)\n",
+    "\n",
+    "        prompt_ids.extend(enc['input_ids'][0].tolist())\n",
+    "        completion_ids.extend(cids)\n",
+    "        logprobs.extend([\n",
+    "            F.log_softmax(s[0], dim=-1)[t].item()\n",
+    "            for s, t in zip(out.scores, cids)\n",
+    "        ])\n",
+    "\n",
+    "        if step == 0:\n",
+    "            print(f'  [sample] {repr(completion_text[:180])}')\n",
+    "\n",
+    "        parsed   = extract_json_action(completion_text)\n",
+    "        is_valid = parsed is not None\n",
+    "        if not is_valid:\n",
+    "            parsed = step_aware_fallback(step, max_steps)\n",
+    "        else:\n",
+    "            valid_json_count += 1\n",
+    "\n",
+    "        action_type = parsed.get('action_type', 'meta.noop')\n",
+    "        args        = parsed.get('args', {})\n",
+    "        print(f'  [step {step}] {action_type}')\n",
+    "\n",
+    "        # Compliance signal capture\n",
+    "        if action_type == 'meta.read_runbook' and is_valid:\n",
+    "            read_runbook_done = True\n",
+    "        if action_type == 'ticketing.create_ticket' and is_valid and ticket_label is None:\n",
+    "            ticket_label    = args.get('label')\n",
+    "            ticket_priority = args.get('priority')\n",
+    "        if action_type == 'ticketing.assign_ticket' and is_valid and assigned_team is None:\n",
+    "            assigned_team = args.get('team')\n",
+    "        if action_type == 'chat.post_message' and is_valid:\n",
+    "            ch = args.get('channel', '')\n",
+    "            if ch:\n",
+    "                posted_channels.append(ch)\n",
+    "\n",
+    "        turn_history.append({\n",
+    "            'obs_text'  : obs_text,\n",
+    "            'completion': completion_text,\n",
+    "            'is_runbook': (action_type == 'meta.read_runbook' and is_valid),\n",
+    "        })\n",
+    "\n",
+    "        result   = env.step({'action_type': action_type, 'args': args})\n",
+    "        obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "\n",
+    "        # Extract org config from runbook response (one step after the call)\n",
+    "        last_result = obs_dict.get('last_action_result') or {}\n",
+    "        if action_type == 'meta.read_runbook' and last_result.get('ok'):\n",
+    "            data = last_result.get('data') or {}\n",
+    "            if isinstance(data, dict):\n",
+    "                org = data.get('org_config') or {}\n",
+    "                valid_labels     = set(org.get('label_taxonomy', {}).values())\n",
+    "                valid_priorities = set(org.get('priority_levels', []))\n",
+    "                valid_teams      = set(org.get('team_map', {}).values())\n",
+    "                oncall_channels  = set(org.get('oncall_channels', {}).values())\n",
+    "\n",
+    "        done      = bool(getattr(result, 'done', obs_dict.get('done', False)))\n",
+    "        env_score = float(getattr(result, 'reward', obs_dict.get('reward', 0.0)))\n",
+    "        step     += 1\n",
+    "\n",
+    "    reward = compute_rollout_reward(\n",
+    "        read_runbook_done = read_runbook_done,\n",
+    "        valid_labels      = valid_labels,\n",
+    "        valid_priorities  = valid_priorities,\n",
+    "        valid_teams       = valid_teams,\n",
+    "        oncall_channels   = oncall_channels,\n",
+    "        ticket_label      = ticket_label,\n",
+    "        ticket_priority   = ticket_priority,\n",
+    "        assigned_team     = assigned_team,\n",
+    "        posted_channels   = posted_channels,\n",
+    "        env_score         = env_score,\n",
+    "        valid_json_count  = valid_json_count,\n",
+    "    )\n",
+    "\n",
+    "    lbl = '✓' if ticket_label    and ticket_label    in valid_labels     else ('✗' if ticket_label    else '-')\n",
+    "    pri = '✓' if ticket_priority and ticket_priority in valid_priorities  else ('✗' if ticket_priority else '-')\n",
+    "    tm  = '✓' if assigned_team   and assigned_team   in valid_teams       else ('✗' if assigned_team   else '-')\n",
+    "    ch  = '✓' if any(c in oncall_channels for c in posted_channels)       else ('✗' if posted_channels else '-')\n",
+    "    print(f'  [rollout] steps={step} env={env_score:.3f} '\n",
+    "          f'label={lbl} priority={pri} team={tm} channel={ch} → reward={reward:.3f}')\n",
+    "\n",
+    "    return {\n",
+    "        'prompt_ids'    : prompt_ids,\n",
+    "        'completion_ids': completion_ids,\n",
+    "        'logprobs'      : logprobs,\n",
+    "        'reward'        : reward,\n",
+    "    }\n",
+    "\n",
+    "\n",
+    "def grpo_rollout_func(prompts, trainer=None):\n",
+    "    out = {'prompt_ids': [], 'completion_ids': [], 'logprobs': [], 'reward': []}\n",
+    "    prompt_seen: dict = {}\n",
+    "    for prompt in prompts:\n",
+    "        gen_offset          = prompt_seen.get(prompt, 0)\n",
+    "        prompt_seen[prompt] = gen_offset + 1\n",
+    "        ep = run_grpo_episode(\n",
+    "            trainer, grpo_env, tokenizer, prompt, TRAIN_MAX_STEPS, gen_slot=gen_offset\n",
+    "        )\n",
+    "        for k in out:\n",
+    "            out[k].append(ep[k])\n",
+    "    return out\n",
+    "\n",
+    "\n",
+    "def grpo_reward_func(completions, **kwargs):\n",
+    "    \"\"\"Passthrough — reward pre-computed in grpo_rollout_func.\"\"\"\n",
+    "    rewards = kwargs.get('reward', [])\n",
+    "    if not rewards:\n",
+    "        print(f'[grpo_reward_func] no reward in kwargs — keys: {list(kwargs.keys())}')\n",
+    "        return [0.0] * len(completions)\n",
+    "    return [float(r) for r in rewards]\n",
+    "\n",
+    "\n",
+    "print(f'GRPO rollout ready  max_steps={TRAIN_MAX_STEPS}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s12",
+   "metadata": {},
+   "source": ["## 12. GRPO Config + Trainer"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-trainer",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from trl import GRPOConfig\n",
+    "from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
+    "\n",
+    "OUTPUT_DIR = 'pm-ops-grpo-Qwen3-1.7B-triage-v4'\n",
+    "HF_REPO_ID = f'Saurav1/{OUTPUT_DIR}'\n",
+    "\n",
+    "grpo_cfg = GRPOConfig(\n",
+    "    num_train_epochs            = 2,\n",
+    "    learning_rate               = 1e-6,\n",
+    "    gradient_accumulation_steps = GRAD_ACCUM,\n",
+    "    per_device_train_batch_size = 1,\n",
+    "    warmup_steps                = 5,\n",
+    "    num_generations             = NUM_GEN,\n",
+    "    max_completion_length       = MAX_COMP_LEN,\n",
+    "    max_prompt_length           = 4096,\n",
+    "    use_vllm                    = False,\n",
+    "    output_dir                  = OUTPUT_DIR,\n",
+    "    report_to                   = 'none',\n",
+    "    logging_steps               = 1,\n",
+    "    save_steps                  = 20,\n",
+    "    gradient_checkpointing      = False,\n",
+    ")\n",
+    "\n",
+    "eff_batch   = grpo_cfg.per_device_train_batch_size * GRAD_ACCUM\n",
+    "total_steps = len(grpo_dataset) * NUM_GEN * grpo_cfg.num_train_epochs // eff_batch\n",
+    "print(f'GRPO: {len(grpo_dataset)} eps × {NUM_GEN} gen × {grpo_cfg.num_train_epochs} epochs → ~{total_steps} steps')\n",
+    "\n",
+    "trainer = PMOpsGRPOTrainer(\n",
+    "    model            = model,\n",
+    "    processing_class = tokenizer,\n",
+    "    reward_funcs     = grpo_reward_func,\n",
+    "    train_dataset    = grpo_dataset,\n",
+    "    args             = grpo_cfg,\n",
+    "    rollout_func     = grpo_rollout_func,\n",
+    ")\n",
+    "assert grpo_cfg.use_vllm is False\n",
+    "print(f'Trainer: {type(trainer).__name__}  ready')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s13",
+   "metadata": {},
+   "source": ["## 13. Preflight Probe"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-probe",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# Run 2 episodes and verify reward varies (not stuck at a constant)\n",
+    "probe = grpo_rollout_func(\n",
+    "    [grpo_dataset[0]['prompt'], grpo_dataset[1]['prompt']],\n",
+    "    trainer=trainer,\n",
+    ")\n",
+    "print(f'probe rewards : {probe[\"reward\"]}')\n",
+    "assert len(probe['reward']) == 2, 'need one reward per prompt'\n",
+    "assert all(isinstance(r, float) for r in probe['reward']), 'rewards must be float'\n",
+    "print('Preflight PASS ✓')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s14",
+   "metadata": {},
+   "source": [
+    "## 14. Train\n",
+    "\n",
+    "Watch for:\n",
+    "- `[sample] '\\`\\`\\`json ...'` — model should output JSON code blocks\n",
+    "- `label=✓ priority=✓ team=✓ channel=✓` — compliance signals the model is getting right\n",
+    "- `reward/injected_std > 0` in logs — confirms GRPO has a non-zero gradient signal"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-train",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "trainer_stats = trainer.train()\n",
+    "\n",
+    "train_mins = round(trainer_stats.metrics.get('train_runtime', 0) / 60, 1)\n",
+    "used_gb    = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
+    "print(f'Training time : {train_mins} min')\n",
+    "print(f'Peak GPU      : {used_gb} GB / {TOTAL_GB} GB ({round(used_gb/TOTAL_GB*100, 1)}%)')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s15",
+   "metadata": {},
+   "source": ["## 15. Save + Push to HF"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-save",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "grpo_env.close()\n",
+    "\n",
+    "model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\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",
+   "id": "v4-eval-hdr",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "## Evaluation — Baseline vs Trained"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s16",
+   "metadata": {},
+   "source": ["## 16. Evaluate"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-eval",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "from training.rollout import extract_json_action, step_aware_fallback\n",
+    "from inference import baseline_agent\n",
+    "\n",
+    "N_EVAL         = 15\n",
+    "EVAL_MAX_STEPS = 12\n",
+    "EVAL_SEED_BASE = 9000\n",
+    "\n",
+    "\n",
+    "def run_eval(n=N_EVAL):\n",
+    "    scores = []\n",
+    "    model.eval()\n",
+    "    with GenericEnvClient(base_url=ENV_URL).sync() as env:\n",
+    "        for i in range(n):\n",
+    "            result     = env.reset(seed=EVAL_SEED_BASE + i)\n",
+    "            obs_dict   = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\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   = tokenizer.apply_chat_template(\n",
+    "                    msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
+    "                )\n",
+    "                inputs = tokenizer([prompt], return_tensors='pt', truncation=True, max_length=4096)\n",
+    "                inputs = {k: v.to(model.device) for k, v in inputs.items()}\n",
+    "                with torch.no_grad():\n",
+    "                    out_ids = model.generate(\n",
+    "                        **inputs, max_new_tokens=256, do_sample=False,\n",
+    "                        pad_token_id=tokenizer.eos_token_id,\n",
+    "                    )\n",
+    "                completion = tokenizer.decode(\n",
+    "                    out_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True\n",
+    "                )\n",
+    "                parsed   = extract_json_action(completion) or step_aware_fallback(step, EVAL_MAX_STEPS)\n",
+    "                result   = env.step({'action_type': parsed['action_type'], 'args': parsed.get('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",
+    "                history.append({'obs_text': obs_text, 'completion': completion, 'is_runbook': False})\n",
+    "                step += 1\n",
+    "\n",
+    "            scores.append(score)\n",
+    "        model.train()\n",
+    "    return scores\n",
+    "\n",
+    "\n",
+    "def run_baseline(n=N_EVAL):\n",
+    "    scores = []\n",
+    "    with GenericEnvClient(base_url=ENV_URL).sync() as env:\n",
+    "        for i in range(n):\n",
+    "            result   = env.reset(seed=EVAL_SEED_BASE + i)\n",
+    "            obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
+    "            org_config, step, score, done = {}, 0, 0.0, False\n",
+    "            while not done and step < EVAL_MAX_STEPS:\n",
+    "                at, args = baseline_agent(obs_dict, org_config)\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",
+    "                if at == 'meta.read_runbook':\n",
+    "                    last = obs_dict.get('last_action_result') or {}\n",
+    "                    if last.get('ok'):\n",
+    "                        data = last.get('data') or {}\n",
+    "                        if isinstance(data, dict) and 'org_config' in data:\n",
+    "                            org_config.update(data['org_config'])\n",
+    "                step += 1\n",
+    "            scores.append(score)\n",
+    "    return scores\n",
+    "\n",
+    "\n",
+    "print('--- Baseline ---')\n",
+    "baseline_scores = run_baseline()\n",
+    "print('--- Trained ---')\n",
+    "trained_scores  = run_eval()\n",
+    "\n",
+    "b_avg = sum(baseline_scores) / N_EVAL\n",
+    "t_avg = sum(trained_scores)  / N_EVAL\n",
+    "print(f'\\nBaseline avg : {b_avg:.3f}')\n",
+    "print(f'Trained  avg : {t_avg:.3f}')\n",
+    "print(f'Delta        : {t_avg - b_avg:+.3f}')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "v4-s17",
+   "metadata": {},
+   "source": ["## 17. Teardown"]
+  },
+  {
+   "cell_type": "code",
+   "id": "v4-teardown",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "server_proc.terminate()\n",
+    "print('PM-Ops server stopped')"
+   ]
+  }
+ ]
+}