{
 "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 | `rollout_func` kwargs broken in TRL 1.3.0.dev0 | TRL **1.2.0** stable + single `reward` key |\n",
    "| Dep hell | Unsloth 2026.4.8 requires torch<2.11 / trl≤0.24, conflicts with Colab's base image | **Drop Unsloth**, use bitsandbytes + PEFT directly |\n",
    "| Only 4 steps | T4 + grad_accum=64 | GPU auto-detect, adaptive config |\n",
    "| Silent reward failure | No logging | Rollout prints every episode, reward_func warns on zero |\n",
    "\n",
    "**Stack:** TRL 1.2.0 · bitsandbytes (4-bit) · PEFT LoRA · OpenEnv  \n",
    "**GPU:** A100 → ~150 steps / 60 min · T4 → ~18 steps (smoke-test / short run)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Keep Colab's native torch (2.11.0) — don't touch it.\n",
    "# Install everything that is compatible with torch 2.11.0.\n",
    "\n",
    "!pip install -q \"trl==1.2.0\"\n",
    "!pip install -q \"bitsandbytes>=0.44.0\"          # 4-bit quantisation\n",
    "!pip install -q \"peft>=0.13.0\"                  # LoRA\n",
    "!pip install -q \"accelerate>=1.0.0\"\n",
    "!pip install -q \"datasets>=4.7.0\"\n",
    "!pip install -q \"transformers>=5.0.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\n",
    "import trl.experimental.openenv  # must exist for rollout_func\n",
    "\n",
    "print(f'torch : {torch.__version__}')\n",
    "print(f'TRL   : {trl.__version__}  (need 1.2.x)')\n",
    "assert trl.__version__.startswith('1.2'), \\\n",
    "    f'Wrong TRL: {trl.__version__}. Re-run install cell and restart kernel.'\n",
    "print('trl.experimental.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'\\nGPU  : {gpu.name}  ({TOTAL_GB} GB)')\n",
    "print(f'Mode : {\"A100 — full training\" if IS_A100 else \"T4 — reduced config (smoke-test quality)\"}')\n",
    "\n",
    "# Adaptive config\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",
    "\n",
    "print(f'\\ngrad_accum={GRAD_ACCUM}  num_gen={NUM_GEN}  max_comp={MAX_COMP_LEN}')"
   ]
  },
  {
   "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",
    "    !git -C {REPO_DIR} pull -q origin main\n",
    "    print(f'Pulled latest → {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()}')"
   ]
  },
  {
   "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",
    "\n",
    "Localhost removes ~200 ms/step network round-trip vs calling the HF Space."
   ]
  },
  {
   "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 _ 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",
    "from training.rollout import _obs_to_dict\n",
    "\n",
    "with GenericEnvClient(base_url=ENV_URL).sync() as _env:\n",
    "    r = _env.reset()\n",
    "    obs = _obs_to_dict(r.observation if hasattr(r, 'observation') else r)\n",
    "    print(f'task_brief : {obs.get(\"task_brief\", \"?\")[:80]}...')\n",
    "    _env.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
    "    print('Env step   : OK')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Load Model — bitsandbytes 4-bit + PEFT LoRA\n",
    "\n",
    "No Unsloth — uses stock HuggingFace `AutoModelForCausalLM` with `BitsAndBytesConfig`.  \n",
    "Same 4-bit + LoRA memory profile, ~2× slower generation than Unsloth's kernels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "MODEL_NAME = 'Qwen/Qwen3-1.7B'\n",
    "LORA_RANK  = 16\n",
    "\n",
    "bnb_cfg = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type='nf4',\n",
    ")\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    MODEL_NAME,\n",
    "    quantization_config=bnb_cfg,\n",
    "    device_map='auto',\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    attn_implementation='eager',  # safe default; use 'flash_attention_2' if flash-attn installed\n",
    ")\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = 'left'\n",
    "\n",
    "lora_cfg = LoraConfig(\n",
    "    r=LORA_RANK,\n",
    "    lora_alpha=LORA_RANK,\n",
    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
    "    lora_dropout=0.0,\n",
    "    bias='none',\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    ")\n",
    "model = get_peft_model(model, lora_cfg)\n",
    "model.print_trainable_parameters()\n",
    "\n",
    "reserved = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
    "print(f'GPU memory after load: {reserved} GB / {TOTAL_GB} GB')"
   ]
  },
  {
   "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",
    "dataset = Dataset.from_list([{'prompt': r['prompt']} for r in rows])\n",
    "print(f'Dataset : {len(dataset)} triage episodes')\n",
    "print(f'Sample  : {dataset[0][\"prompt\"][:100]}...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Create Persistent Env Client + Rollout Function\n",
    "\n",
    "**Key change from v1:** single `\"reward\"` key in the rollout output.  \n",
    "This removes all kwargs-plumbing complexity — the reward_func is a trivial passthrough."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openenv.core import GenericEnvClient\n",
    "from training.rollout import rollout_once\n",
    "\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",
    "\n",
    "def make_rollout_func(env, tok, max_steps=TRAIN_MAX_STEPS):\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",
    "            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",
    "\n",
    "rollout_func = make_rollout_func(sync_env, tokenizer)\n",
    "\n",
    "\n",
    "def reward_func(completions, **kwargs):\n",
    "    \"\"\"Trivial passthrough — reads pre-computed reward from rollout kwargs.\"\"\"\n",
    "    n       = len(completions)\n",
    "    rewards = kwargs.get('reward', [])\n",
    "    if not rewards:\n",
    "        print(\n",
    "            f'[ERROR] reward_func: empty kwargs! '\n",
    "            f'TRL {trl.__version__} is not passing rollout keys. '\n",
    "            f'kwargs keys present: {list(kwargs.keys())}'\n",
    "        )\n",
    "        return [0.0] * n\n",
    "    return [float(r) for r in rewards]\n",
    "\n",
    "\n",
    "print('rollout_func + reward_func defined')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Smoke Test\n",
    "\n",
    "Verify the env + reward pipeline works *before* committing to a 60-minute training run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('─── Smoke test ──────────────────────────────────')\n",
    "\n",
    "# 1. Env resets + steps correctly\n",
    "with GenericEnvClient(base_url=ENV_URL).sync() as _t:\n",
    "    r1 = _t.reset(seed=42)\n",
    "    obs = _obs_to_dict(r1.observation if hasattr(r1, 'observation') else r1)\n",
    "    assert obs.get('task_brief'), 'FAIL: no task_brief in observation'\n",
    "    print(f'[OK] env.reset()       task_brief present')\n",
    "\n",
    "    r2 = _t.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
    "    print(f'[OK] env.step(read_runbook)')\n",
    "\n",
    "    r3 = _t.step({'action_type': 'meta.finish', 'args': {}})\n",
    "    obs3 = _obs_to_dict(r3.observation if hasattr(r3, 'observation') else r3)\n",
    "    final = float(getattr(r3, 'reward', obs3.get('reward', -1)))\n",
    "    print(f'[OK] env.step(finish)  reward={final:.3f}')\n",
    "\n",
    "# 2. reward_func passthrough\n",
    "rf = reward_func(['c1', 'c2'], reward=[0.3, 0.7])\n",
    "assert rf == [0.3, 0.7], f'FAIL: passthrough broken: {rf}'\n",
    "print(f'[OK] reward_func passthrough  {rf}')\n",
    "\n",
    "# 3. reward_func warns correctly on empty kwargs\n",
    "import io, warnings\n",
    "rf_empty = reward_func(['c1'])\n",
    "assert rf_empty == [0.0]\n",
    "print(f'[OK] reward_func empty → [0.0]  (error printed above is expected)')\n",
    "\n",
    "print('\\nSmoke test PASSED')"
   ]
  },
  {
   "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",
    "    # No vLLM — standard HF generate (compatible with bitsandbytes + PEFT)\n",
    "    use_vllm=False,\n",
    "\n",
    "    # Output\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=True,  # saves ~30% activation memory\n",
    ")\n",
    "\n",
    "eff_batch = grpo_config.per_device_train_batch_size * GRAD_ACCUM\n",
    "total_steps = len(dataset) * NUM_GEN // eff_batch\n",
    "print(f'Output dir   : {OUTPUT_DIR}')\n",
    "print(f'Eff. batch   : {eff_batch}')\n",
    "print(f'Total steps  : ~{total_steps}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Create Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from training.pm_ops_trainer import PMOpsGRPOTrainer\n",
    "\n",
    "trainer = PMOpsGRPOTrainer(\n",
    "    model=model,\n",
    "    processing_class=tokenizer,\n",
    "    reward_funcs=reward_func,   # fallback only — injected rewards take priority\n",
    "    train_dataset=dataset,\n",
    "    args=grpo_config,\n",
    "    rollout_func=rollout_func,\n",
    ")\n",
    "print(\"PMOpsGRPOTrainer ready — direct reward injection active\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. Train\n",
    "\n",
    "**Watch Colab stdout for rollout lines:**\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",
    "- `combined > 0` in stdout **and** `train/reward > 0` in trackio → working correctly  \n",
    "- `combined > 0` in stdout **but** `train/reward = 0` in trackio → kwargs still not flowing  \n",
    "  → Add `print(kwargs.keys())` inside `reward_func` to debug  \n",
    "- `combined = 0` always → check env connection / reward logic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer_stats = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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      : {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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sync_env.close()\n",
    "\n",
    "# Save LoRA adapter (not merged — push separately if needed)\n",
    "trainer.save_model(OUTPUT_DIR)\n",
    "tokenizer.save_pretrained(OUTPUT_DIR)\n",
    "print(f'Saved LoRA adapter → {OUTPUT_DIR}')\n",
    "\n",
    "# Push to Hub\n",
    "trainer.model.push_to_hub(HF_REPO_ID)\n",
    "tokenizer.push_to_hub(HF_REPO_ID)\n",
    "print(f'Pushed → https://huggingface.co/{HF_REPO_ID}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 14. Merge LoRA → bf16 (Optional — for full-weight inference)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Only run this if you have enough CPU RAM (~7 GB free) to hold the merged model.\n",
    "# On T4 Colab this often OOMs — skip and load with PEFT for eval instead.\n",
    "from peft import PeftModel\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "base = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map='cpu')\n",
    "merged = PeftModel.from_pretrained(base, OUTPUT_DIR)\n",
    "merged = merged.merge_and_unload()\n",
    "merged.save_pretrained(f'{OUTPUT_DIR}-merged')\n",
    "tokenizer.save_pretrained(f'{OUTPUT_DIR}-merged')\n",
    "print(f'Merged model saved → {OUTPUT_DIR}-merged')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 15. Evaluate: Baseline vs Trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from training.rollout import extract_json_action, step_aware_fallback, build_messages\n",
    "from training.rollout import _current_obs_text, _obs_to_dict\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",
    "# Put model in eval/inference mode\n",
    "model.eval()\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_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_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', truncation=True, max_length=4096)\n",
    "            inputs = {k: v.to(mdl.device) for k, v in inputs.items()}\n",
    "            with torch.no_grad():\n",
    "                out_ids    = mdl.generate(**inputs, max_new_tokens=512,\n",
    "                                          do_sample=False, pad_token_id=tok.eos_token_id)\n",
    "            completion = tok.decode(out_ids[0][inputs['input_ids'].shape[1]:], 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_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",
    "        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": [
    "## 16. Plot"
   ]
  },
  {
   "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', ls='--', alpha=0.5, lw=1)\n",
    "ax.axhline(sum(trained_scores)/N_EVAL,  color='coral',     ls='--', alpha=0.5, lw=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'Avg 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": [
    "## 17. 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
}