{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PM-Ops: Train a Project Management Agent with GRPO + Unsloth\n",
    "\n",
    "Fine-tune **Qwen3-1.7B** on PM-Ops triage tasks using GRPO (TRL) + Unsloth for memory efficiency.\n",
    "\n",
    "**GPU:** A100 40GB (HF Jupyter Space — uses HF credits)  \n",
    "**Time:** ~90 min (150 triage episodes, 1 epoch)  \n",
    "**Stack:** OpenEnv + TRL + Unsloth (hackathon-recommended stack)\n",
    "\n",
    "### Architecture\n",
    "```\n",
    "[HF A100 GPU Space]\n",
    "  ├── Unsloth + GRPOTrainer  ← Qwen3-1.7B with LoRA, 4-bit quantised\n",
    "  └── PM-Ops FastAPI (localhost:8000)  ← subprocess, <1ms per step\n",
    "```\n",
    "\n",
    "### Reward design (5 independent signals — anti-hack per hackathon guide)\n",
    "| Signal | Weight | Purpose |\n",
    "|---|---|---|\n",
    "| `reward_final_score` | 0.45 | Correctness: label + priority + team + channel |\n",
    "| `reward_no_wrong_channels` | 0.15 | Anti-hack: penalise channel-spray |\n",
    "| `reward_valid_json` | 0.15 | Format discipline |\n",
    "| `reward_read_runbook` | 0.15 | Process: read before acting |\n",
    "| `reward_efficiency` | 0.10 | Speed when correct |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unsloth first — it pins compatible versions of torch/transformers\n",
    "!pip install -Uq unsloth\n",
    "!pip install -Uq \"trl>=0.17.0\" openenv-core datasets trackio\n",
    "print('Dependencies installed.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Clone PM-Ops Repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, sys\n",
    "\n",
    "REPO_URL = 'https://huggingface.co/spaces/adityaguntur/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",
    "    print(f'Already exists: {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'Working directory: {os.getcwd()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. 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": [
    "## 3. Start PM-Ops Server Locally\n",
    "\n",
    "Running server on localhost eliminates the ~200ms/step network cost of calling the HF Space.\n",
    "With 15 steps × 150 episodes × 2 generations that's 4,500 saved round-trips."
   ]
  },
  {
   "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',\n",
    "     '--host', '0.0.0.0', '--port', '8000'],\n",
    "    cwd=REPO_DIR,\n",
    "    stdout=subprocess.DEVNULL,\n",
    "    stderr=subprocess.DEVNULL,\n",
    ")\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 at {ENV_URL}  (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 — check uvicorn install.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. 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 _check:\n",
    "    res = _check.reset()\n",
    "    obs = res.observation if hasattr(res, 'observation') else res\n",
    "    brief = getattr(obs, 'task_brief', '') or obs.get('task_brief', '')\n",
    "    print(f'Task brief: {brief[:120]}...')\n",
    "    res2 = _check.step({'action_type': 'meta.read_runbook', 'args': {}})\n",
    "    print('Step OK — environment is working.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Load Model with Unsloth\n",
    "\n",
    "Unsloth applies:\n",
    "- **4-bit quantisation** — halves GPU memory vs bf16\n",
    "- **LoRA adapters** — only trains ~1% of parameters, much faster\n",
    "- **Unsloth kernel optimisations** — 2× faster rollout generation\n",
    "- **`use_gradient_checkpointing=\"unsloth\"`** — 30% less activation memory\n",
    "\n",
    "> **Save warning (hackathon guide point 16):** never merge LoRA into a 4-bit model directly.\n",
    "> Use `model.save_pretrained_merged(..., save_method=\"merged_16bit\")` in cell 14."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from unsloth import FastLanguageModel, PatchFastRL\n",
    "from trl import GRPOTrainer  # import before patching so patch takes effect\n",
    "\n",
    "PatchFastRL('GRPO', FastLanguageModel)\n",
    "\n",
    "MODEL_NAME  = 'Qwen/Qwen3-1.7B'\n",
    "MAX_SEQ_LEN = 4096 + 512   # max_prompt_length + max_completion_length\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=True,          # enables Unsloth's vLLM-compatible fast path\n",
    "    max_lora_rank=LORA_RANK,\n",
    "    gpu_memory_utilization=0.6,   # ~24 GB on A100 40GB for KV cache\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',\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",
    "print(f'Model ready: {MODEL_NAME} (4-bit + LoRA r={LORA_RANK})')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Generate Training Dataset\n",
    "\n",
    "150 fixed-seed triage episodes. Seed is embedded in the prompt string so `env.reset(seed=...)` \n",
    "reproduces the exact same org config — training is fully reproducible."
   ]
  },
  {
   "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",
    "\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'Difficulties: {set(r[\"difficulty\"] for r in rows)}')\n",
    "print(f'Sample: {dataset[0][\"prompt\"][:120]}...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Create Persistent Environment Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openenv.core import GenericEnvClient\n",
    "\n",
    "sync_env = GenericEnvClient(base_url=ENV_URL).sync()\n",
    "sync_env.connect()\n",
    "print('Persistent training connection established.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Build Rollout Function\n",
    "\n",
    "- `max_steps=15` — triage can be solved in 5 steps; 15 gives exploration room without burning compute  \n",
    "- Runbook-pinned truncation — org config stays in context even when older turns scroll out  \n",
    "- Step-aware fallback — `read_runbook` (early) → `noop` (mid) → `finish` (late)  \n",
    "- Channel tracking — feeds `reward_no_wrong_channels` to catch spray behaviour"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from training.rollout import make_rollout_func\n",
    "\n",
    "TRAIN_MAX_STEPS = 15  # triage; increase for other tasks\n",
    "\n",
    "rollout_func = make_rollout_func(\n",
    "    sync_env=sync_env,\n",
    "    tokenizer=tokenizer,\n",
    "    max_steps=TRAIN_MAX_STEPS,\n",
    ")\n",
    "print(f'Rollout ready (max_steps={TRAIN_MAX_STEPS})')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Define Reward Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from training.rewards import ALL_REWARD_FUNCS, WEIGHT_FINAL_SCORE, WEIGHT_NO_WRONG_CHANNELS\n",
    "print(f'Reward functions ({len(ALL_REWARD_FUNCS)}):')\n",
    "for f in ALL_REWARD_FUNCS:\n",
    "    print(f'  {f.__name__}')"
   ]
  },
  {
   "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'\n",
    "HF_REPO_ID = f'your-hf-username/{OUTPUT_DIR}'  # ← update before running\n",
    "\n",
    "grpo_config = GRPOConfig(\n",
    "    # Training\n",
    "    num_train_epochs=1,\n",
    "    learning_rate=5e-6,\n",
    "    gradient_accumulation_steps=64,\n",
    "    per_device_train_batch_size=1,\n",
    "    warmup_steps=10,\n",
    "    num_generations=2,\n",
    "\n",
    "    # Sequence lengths — longer than Wordle due to JSON + reasoning\n",
    "    max_completion_length=512,\n",
    "    max_prompt_length=4096,\n",
    "\n",
    "    # vLLM — Unsloth handles colocate mode via fast_inference=True above\n",
    "    use_vllm=True,\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",
    "print(f'Output: {OUTPUT_DIR}')\n",
    "print(f'Effective batch size: {grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}')"
   ]
  },
  {
   "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,              # Unsloth model object (not string)\n",
    "    processing_class=tokenizer,\n",
    "    reward_funcs=ALL_REWARD_FUNCS,\n",
    "    train_dataset=dataset,\n",
    "    args=grpo_config,\n",
    "    rollout_func=rollout_func,\n",
    ")\n",
    "print('GRPOTrainer ready.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 12. GPU Check"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "gpu = torch.cuda.get_device_properties(0)\n",
    "reserved_before = round(torch.cuda.max_memory_reserved() / 1024**3, 2)\n",
    "total_gb = round(gpu.total_memory / 1024**3, 2)\n",
    "print(f'GPU: {gpu.name}')\n",
    "print(f'Memory: {total_gb} GB total, {reserved_before} GB reserved')\n",
    "assert total_gb >= 38, f'Need A100 40GB, got {total_gb:.1f} GB — switch runtime.'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 13. Train (~90 min on A100)\n",
    "\n",
    "Watch **trackio** for per-reward-signal curves. Key signals to monitor:\n",
    "- `reward_final_score` — should trend upward over training\n",
    "- `reward_valid_json` — should stay > 0.7 (model reliably outputs JSON)\n",
    "- `reward_read_runbook` — should hit 1.0 quickly and stay there\n",
    "- `reward_no_wrong_channels` — should increase (less channel spray)\n",
    "\n",
    "If `reward_valid_json` < 0.3 for first 20 steps → model struggling with format. Stop and reduce `max_completion_length` to 256."
   ]
  },
  {
   "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 usage: {used_gb} GB / {total_gb} GB ({round(used_gb/total_gb*100, 1)}%)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 14. Save Model\n",
    "\n",
    "**Important:** Unsloth uses LoRA + 4-bit quantisation. Do NOT use `trainer.save_model()` directly —\n",
    "merging LoRA into a 4-bit base corrupts weights (hackathon guide point 16).\n",
    "Use `save_pretrained_merged` which dequantises first, then merges cleanly into bf16."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sync_env.close()  # close training connection before saving\n",
    "\n",
    "# Merge LoRA into bf16 weights and save — safe Unsloth path\n",
    "model.save_pretrained_merged(OUTPUT_DIR, tokenizer, save_method='merged_16bit')\n",
    "print(f'Model saved → {OUTPUT_DIR}')\n",
    "\n",
    "# Push merged model to HF Hub\n",
    "model.push_to_hub_merged(HF_REPO_ID, tokenizer, save_method='merged_16bit')\n",
    "print(f'Pushed → {HF_REPO_ID}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 15. Evaluate: Baseline vs Trained\n",
    "\n",
    "10 fresh triage episodes on the remote HF Space (not localhost) — clean eval setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM\n",
    "from openenv.core import GenericEnvClient\n",
    "from training.rollout import extract_json_action, step_aware_fallback, build_messages, _obs_to_dict\n",
    "from training.prompts import format_observation\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",
    "eval_model = AutoModelForCausalLM.from_pretrained(\n",
    "    OUTPUT_DIR, torch_dtype='auto', device_map='auto'\n",
    ")\n",
    "\n",
    "\n",
    "def _get(o, k, default=None):\n",
    "    return getattr(o, k, default) if not isinstance(o, dict) else o.get(k, default)\n",
    "\n",
    "\n",
    "def eval_trained(sync_env, model, tokenizer, n=N_EVAL):\n",
    "    scores = []\n",
    "    for i in range(n):\n",
    "        result = sync_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 = [], 0, 0.0\n",
    "        done = False\n",
    "\n",
    "        while not done and step < EVAL_MAX_STEPS:\n",
    "            from training.rollout import _current_obs_text\n",
    "            obs_text = _current_obs_text(obs_dict, step, task_brief)\n",
    "            msgs = build_messages(history, obs_text)\n",
    "            prompt_text = tokenizer.apply_chat_template(\n",
    "                msgs, add_generation_prompt=True, tokenize=False, enable_thinking=False\n",
    "            )\n",
    "            inputs = tokenizer([prompt_text], return_tensors='pt').to(model.device)\n",
    "            out_ids = model.generate(**inputs, max_new_tokens=512)\n",
    "            completion = tokenizer.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 = sync_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(sync_env, n=N_EVAL):\n",
    "    scores = []\n",
    "    for i in range(n):\n",
    "        result = sync_env.reset()\n",
    "        obs_dict = _obs_to_dict(result.observation if hasattr(result, 'observation') else result)\n",
    "        org_config, done, step, score = {}, False, 0, 0.0\n",
    "        while not done and step < EVAL_MAX_STEPS:\n",
    "            at, args = baseline_agent(obs_dict, org_config)\n",
    "            result = sync_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, 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 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 = axes[0]\n",
    "x, w = np.arange(N_EVAL), 0.35\n",
    "ax.bar(x - w/2, baseline_scores, w, label='Baseline (heuristic)', color='steelblue', alpha=0.8)\n",
    "ax.bar(x + w/2, trained_scores, w, label='Trained (GRPO+Unsloth)', 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('Eval episode')\n",
    "ax.set_ylabel('Episode reward (0–1)')\n",
    "ax.set_title('Per-episode: Baseline vs GRPO-trained')\n",
    "ax.set_xticks(x)\n",
    "ax.set_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+Unsloth'], 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 (0–1)')\n",
    "ax2.set_title(f'Average over {N_EVAL} triage episodes')\n",
    "ax2.set_ylim(0, 1.05)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('eval_results.png', dpi=150, bbox_inches='tight')\n",
    "plt.show()\n",
    "print('Saved: eval_results.png  ← embed this in README for hackathon submission')"
   ]
  },
  {
   "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.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "## What to try next\n",
    "\n",
    "| Change | Where | Expected effect |\n",
    "|---|---|---|\n",
    "| All 4 task types | `generate_triage_dataset` → all-tasks generator | Tests generalization |\n",
    "| 500 episodes | `N_EPISODES = 500` | Better coverage |\n",
    "| Equal weights (0.20 each) | `WEIGHT_*` in `rewards.py` | Compare reward-shaping approaches |\n",
    "| Longer training | `num_train_epochs=3` | More improvement |\n",
    "| Larger model | `Qwen/Qwen3-4B` | Better reasoning, more memory |\n",
    "| Higher LoRA rank | `r=64, lora_alpha=64` | More capacity |\n",
    "| More steps for other tasks | `TRAIN_MAX_STEPS = 20` | Needed for release_notes/dep_update |"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.0"
  },
  "accelerator": "GPU",
  "gpuClass": "premium"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}