Remove redundant train_grpo_safe.ipynb

train_grpo_safe.ipynb

@@ -1,562 +0,0 @@
-{
- "nbformat": 4,
- "nbformat_minor": 0,
- "metadata": {
-  "colab": {
-   "provenance": [],
-   "gpuType": "T4"
-  },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python"
-  },
-  "accelerator": "GPU"
- },
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Support Ticket Env - GRPO Fine-Tuning\n",
-    "**OpenEnv x Scalar Hackathon**\n",
-    "\n",
-    "Fine-tunes `Qwen/Qwen2.5-0.5B-Instruct` using GRPO (Group Relative Policy Optimization) from HuggingFace TRL against the live Support Ticket Environment API.\n",
-    "\n",
-    "- Model: Qwen2.5-0.5B-Instruct\n",
-    "- Algorithm: GRPO\n",
-    "- Environment: https://algocore-support-ticket-env.hf.space\n",
-    "- Runtime: ~45-60 min on free Colab T4"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install -q trl transformers peft accelerate\n",
-    "!pip install -q torch bitsandbytes requests datasets\n",
-    "print('Done')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "HF_TOKEN     = \"YOUR_HF_TOKEN_HERE\"\n",
-    "ENV_BASE_URL = \"https://algocore-support-ticket-env.hf.space\"\n",
-    "MODEL_NAME   = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
-    "OUTPUT_DIR   = \"/content/support-ticket-grpo\"\n",
-    "HF_REPO_ID   = \"AlgoCore/support-ticket-grpo-model\"\n",
-    "\n",
-    "os.environ[\"HF_TOKEN\"] = HF_TOKEN\n",
-    "os.environ[\"HUGGING_FACE_HUB_TOKEN\"] = HF_TOKEN\n",
-    "\n",
-    "import torch\n",
-    "print(\"GPU:\", torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"NO GPU - switch runtime!\")\n",
-    "if torch.cuda.is_available():\n",
-    "    print(\"VRAM:\", round(torch.cuda.get_device_properties(0).total_memory / 1e9, 1), \"GB\")\n",
-    "print(\"Model:\", MODEL_NAME)\n",
-    "print(\"Env:\", ENV_BASE_URL)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import requests\n",
-    "import json\n",
-    "import re\n",
-    "from dataclasses import dataclass\n",
-    "from typing import Optional\n",
-    "\n",
-    "@dataclass\n",
-    "class Obs:\n",
-    "    ticket_id: str\n",
-    "    ticket_text: str\n",
-    "    task_id: int\n",
-    "    current_category: Optional[str]\n",
-    "    resolved: bool\n",
-    "    step_count: int\n",
-    "    feedback: str\n",
-    "    score: float\n",
-    "    reward: float\n",
-    "    done: bool\n",
-    "\n",
-    "class SupportEnvClient:\n",
-    "    def __init__(self, base_url):\n",
-    "        self.base_url = base_url.rstrip('/')\n",
-    "        self.session = requests.Session()\n",
-    "        self.session.headers.update({'Content-Type': 'application/json'})\n",
-    "\n",
-    "    def health(self):\n",
-    "        try:\n",
-    "            r = self.session.get(f\"{self.base_url}/health\", timeout=10)\n",
-    "            return r.status_code == 200\n",
-    "        except:\n",
-    "            return False\n",
-    "\n",
-    "    def reset(self, task_id=1, seed=42):\n",
-    "        r = self.session.post(f\"{self.base_url}/reset\", json={\"task_id\": task_id, \"seed\": seed}, timeout=15)\n",
-    "        r.raise_for_status()\n",
-    "        return self._parse(r.json())\n",
-    "\n",
-    "    def step(self, action):\n",
-    "        r = self.session.post(f\"{self.base_url}/step\", json={\"action\": action}, timeout=15)\n",
-    "        r.raise_for_status()\n",
-    "        return self._parse(r.json())\n",
-    "\n",
-    "    def _parse(self, data):\n",
-    "        obs = data.get('observation', data)\n",
-    "        return Obs(\n",
-    "            ticket_id=obs.get('ticket_id', ''),\n",
-    "            ticket_text=obs.get('ticket_text', ''),\n",
-    "            task_id=obs.get('task_id', 1),\n",
-    "            current_category=obs.get('current_category'),\n",
-    "            resolved=obs.get('resolved', False),\n",
-    "            step_count=obs.get('step_count', 0),\n",
-    "            feedback=obs.get('feedback', ''),\n",
-    "            score=obs.get('score', 0.0),\n",
-    "            reward=obs.get('reward', 0.0),\n",
-    "            done=obs.get('done', False),\n",
-    "        )\n",
-    "\n",
-    "env_client = SupportEnvClient(ENV_BASE_URL)\n",
-    "if env_client.health():\n",
-    "    print('Environment API reachable')\n",
-    "    obs = env_client.reset(task_id=1, seed=42)\n",
-    "    print(f'Ticket: {obs.ticket_id} - {obs.ticket_text[:70]}')\n",
-    "else:\n",
-    "    print('Cannot reach environment - check ENV_BASE_URL')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
-    "import torch\n",
-    "\n",
-    "print(f\"Loading {MODEL_NAME}...\")\n",
-    "\n",
-    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN, trust_remote_code=True)\n",
-    "tokenizer.pad_token = tokenizer.eos_token\n",
-    "tokenizer.padding_side = 'left'\n",
-    "\n",
-    "model = AutoModelForCausalLM.from_pretrained(\n",
-    "    MODEL_NAME,\n",
-    "    token=HF_TOKEN,\n",
-    "    torch_dtype=torch.float16,\n",
-    "    device_map='auto',\n",
-    "    trust_remote_code=True,\n",
-    ")\n",
-    "\n",
-    "print(f'Model loaded - {sum(p.numel() for p in model.parameters())/1e6:.0f}M parameters')\n",
-    "print(f'Device: {next(model.parameters()).device}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from peft import LoraConfig, get_peft_model, TaskType\n",
-    "\n",
-    "lora_config = LoraConfig(\n",
-    "    task_type=TaskType.CAUSAL_LM,\n",
-    "    r=16,\n",
-    "    lora_alpha=32,\n",
-    "    lora_dropout=0.05,\n",
-    "    target_modules=[\"q_proj\", \"v_proj\", \"k_proj\", \"o_proj\"],\n",
-    "    bias=\"none\",\n",
-    ")\n",
-    "\n",
-    "model = get_peft_model(model, lora_config)\n",
-    "model.print_trainable_parameters()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "SYSTEM_PROMPT = \"\"\"You are a customer support AI agent. Given a ticket, respond with a JSON action.\n",
-    "\n",
-    "Respond ONLY with valid JSON:\n",
-    "{\"action_type\": \"classify\"|\"reply\"|\"escalate\"|\"close\", \"category\": \"billing\"|\"technical\"|\"account\"|\"general\"|\"refund\", \"reply_text\": \"...\", \"reason\": \"...\"}\n",
-    "\n",
-    "Rules:\n",
-    "- Task 1: action_type=classify, pick correct category\n",
-    "- Task 2: first classify, then reply/escalate/close\n",
-    "- Task 3: classify each ticket then resolve it\n",
-    "- category only needed for classify\n",
-    "- reply_text only needed for reply\n",
-    "- technical issues: escalate\n",
-    "- resolved issues: close\n",
-    "- billing/account/refund: reply\"\"\"\n",
-    "\n",
-    "def build_prompt(obs):\n",
-    "    user_msg = json.dumps({\n",
-    "        \"ticket_id\": obs.ticket_id,\n",
-    "        \"ticket_text\": obs.ticket_text,\n",
-    "        \"task_id\": obs.task_id,\n",
-    "        \"current_category\": obs.current_category,\n",
-    "        \"feedback\": obs.feedback,\n",
-    "        \"step_count\": obs.step_count,\n",
-    "    }, indent=2)\n",
-    "    messages = [\n",
-    "        {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "        {\"role\": \"user\", \"content\": user_msg},\n",
-    "    ]\n",
-    "    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
-    "\n",
-    "def parse_action(text):\n",
-    "    text = text.strip()\n",
-    "    text = re.sub(r'^```(?:json)?\\s*', '', text)\n",
-    "    text = re.sub(r'\\s*```$', '', text)\n",
-    "    try:\n",
-    "        return json.loads(text)\n",
-    "    except:\n",
-    "        match = re.search(r'\\{.*?\\}', text, re.DOTALL)\n",
-    "        if match:\n",
-    "            try:\n",
-    "                return json.loads(match.group())\n",
-    "            except:\n",
-    "                pass\n",
-    "    return {\"action_type\": \"classify\", \"category\": \"general\"}\n",
-    "\n",
-    "obs = env_client.reset(task_id=1, seed=42)\n",
-    "prompt = build_prompt(obs)\n",
-    "print('Prompt builder OK')\n",
-    "print(f'Prompt length: {len(prompt)} chars')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import random\n",
-    "\n",
-    "SEEDS    = [42, 7, 123, 0, 99]\n",
-    "TASK_IDS = [1, 2, 3]\n",
-    "MAX_STEPS = 6\n",
-    "\n",
-    "def generate_action(prompt, max_new_tokens=150):\n",
-    "    inputs = tokenizer(prompt, return_tensors='pt', truncation=True, max_length=1024).to(model.device)\n",
-    "    with torch.no_grad():\n",
-    "        outputs = model.generate(\n",
-    "            **inputs,\n",
-    "            max_new_tokens=max_new_tokens,\n",
-    "            do_sample=True,\n",
-    "            temperature=0.7,\n",
-    "            top_p=0.9,\n",
-    "            pad_token_id=tokenizer.eos_token_id,\n",
-    "        )\n",
-    "    new_tokens = outputs[0][inputs['input_ids'].shape[1]:]\n",
-    "    return tokenizer.decode(new_tokens, skip_special_tokens=True)\n",
-    "\n",
-    "def run_episode(task_id, seed):\n",
-    "    obs = env_client.reset(task_id=task_id, seed=seed)\n",
-    "    prompts, completions, rewards = [], [], []\n",
-    "    for _ in range(MAX_STEPS):\n",
-    "        if obs.done:\n",
-    "            break\n",
-    "        prompt = build_prompt(obs)\n",
-    "        completion = generate_action(prompt)\n",
-    "        action = parse_action(completion)\n",
-    "        try:\n",
-    "            obs = env_client.step(action)\n",
-    "            reward = float(obs.reward or 0.0)\n",
-    "        except:\n",
-    "            reward = -0.1\n",
-    "            obs.done = True\n",
-    "        prompts.append(prompt)\n",
-    "        completions.append(completion)\n",
-    "        rewards.append(reward)\n",
-    "        if obs.done:\n",
-    "            break\n",
-    "    return prompts, completions, sum(rewards)\n",
-    "\n",
-    "print('Running smoke test...')\n",
-    "p, c, r = run_episode(task_id=1, seed=42)\n",
-    "print(f'Smoke test passed - steps={len(p)}, total_reward={r:.3f}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def evaluate(n_seeds=3):\n",
-    "    results = {}\n",
-    "    seeds = SEEDS[:n_seeds]\n",
-    "    for task_id in [1, 2, 3]:\n",
-    "        task_rewards = []\n",
-    "        for seed in seeds:\n",
-    "            _, _, total = run_episode(task_id=task_id, seed=seed)\n",
-    "            normalized = round(max(0, min(1, total / MAX_STEPS)), 3)\n",
-    "            task_rewards.append(normalized)\n",
-    "        avg = round(sum(task_rewards) / len(task_rewards), 3)\n",
-    "        results[f'task{task_id}'] = avg\n",
-    "        print(f'  Task {task_id}: {avg:.3f}')\n",
-    "    results['overall'] = round(sum(results.values()) / 3, 3)\n",
-    "    print(f'  Overall: {results[\"overall\"]:.3f}')\n",
-    "    return results\n",
-    "\n",
-    "print('=== BASELINE (before training) ===')\n",
-    "baseline_scores = evaluate(n_seeds=3)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from torch.optim import AdamW\n",
-    "from transformers import get_linear_schedule_with_warmup\n",
-    "import numpy as np\n",
-    "\n",
-    "LEARNING_RATE = 5e-5\n",
-    "N_EPISODES    = 60\n",
-    "GROUP_SIZE    = 4\n",
-    "KL_COEFF      = 0.01\n",
-    "GRAD_CLIP     = 1.0\n",
-    "LOG_EVERY     = 5\n",
-    "\n",
-    "optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)\n",
-    "scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=5, num_training_steps=N_EPISODES)\n",
-    "\n",
-    "training_log = []\n",
-    "\n",
-    "print(f'Starting GRPO training: {N_EPISODES} episodes, group_size={GROUP_SIZE}')\n",
-    "print('=' * 60)\n",
-    "\n",
-    "model.train()\n",
-    "\n",
-    "for episode in range(1, N_EPISODES + 1):\n",
-    "    task_id = random.choice(TASK_IDS)\n",
-    "    seed    = random.choice(SEEDS)\n",
-    "\n",
-    "    group_rewards     = []\n",
-    "    group_prompts     = []\n",
-    "    group_completions = []\n",
-    "\n",
-    "    for g in range(GROUP_SIZE):\n",
-    "        obs = env_client.reset(task_id=task_id, seed=seed)\n",
-    "        prompt = build_prompt(obs)\n",
-    "        completion = generate_action(prompt)\n",
-    "        action = parse_action(completion)\n",
-    "        try:\n",
-    "            obs = env_client.step(action)\n",
-    "            reward = float(obs.reward or 0.0)\n",
-    "        except:\n",
-    "            reward = -0.1\n",
-    "        group_rewards.append(reward)\n",
-    "        group_prompts.append(prompt)\n",
-    "        group_completions.append(completion)\n",
-    "\n",
-    "    rewards_arr = np.array(group_rewards, dtype=np.float32)\n",
-    "    advantages  = (rewards_arr - rewards_arr.mean()) / (rewards_arr.std() + 1e-8)\n",
-    "\n",
-    "    total_loss = torch.tensor(0.0, requires_grad=True, device=model.device)\n",
-    "    optimizer.zero_grad()\n",
-    "\n",
-    "    for prompt, completion, adv in zip(group_prompts, group_completions, advantages):\n",
-    "        if not completion.strip():\n",
-    "            continue\n",
-    "        full_text = prompt + completion\n",
-    "        inputs = tokenizer(full_text, return_tensors='pt', truncation=True, max_length=1200).to(model.device)\n",
-    "        prompt_len = tokenizer(prompt, return_tensors='pt')[\"input_ids\"].shape[1]\n",
-    "        outputs = model(**inputs, labels=inputs['input_ids'])\n",
-    "        logits = outputs.logits[:, prompt_len-1:-1, :]\n",
-    "        target_ids = inputs['input_ids'][:, prompt_len:]\n",
-    "        if target_ids.shape[1] == 0:\n",
-    "            continue\n",
-    "        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)\n",
-    "        token_log_probs = log_probs.gather(2, target_ids.unsqueeze(-1)).squeeze(-1)\n",
-    "        seq_log_prob = token_log_probs.mean()\n",
-    "        pg_loss  = -torch.tensor(float(adv), device=model.device) * seq_log_prob\n",
-    "        kl_loss  = KL_COEFF * (seq_log_prob ** 2)\n",
-    "        total_loss = total_loss + (pg_loss + kl_loss) / GROUP_SIZE\n",
-    "\n",
-    "    total_loss.backward()\n",
-    "    torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)\n",
-    "    optimizer.step()\n",
-    "    scheduler.step()\n",
-    "\n",
-    "    avg_reward = float(rewards_arr.mean())\n",
-    "    training_log.append((episode, task_id, avg_reward))\n",
-    "\n",
-    "    if episode % LOG_EVERY == 0:\n",
-    "        print(f'Episode {episode:3d}/{N_EPISODES} | task={task_id} | avg_reward={avg_reward:.3f} | loss={total_loss.item():.4f}')\n",
-    "\n",
-    "print('Training complete!')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.eval()\n",
-    "\n",
-    "print('=== POST-TRAINING EVALUATION ===')\n",
-    "trained_scores = evaluate(n_seeds=3)\n",
-    "\n",
-    "print('\\n=== IMPROVEMENT SUMMARY ===')\n",
-    "print(f'{\"Task\":<10} {\"Before\":>8} {\"After\":>8} {\"Delta\":>8}')\n",
-    "print('-' * 38)\n",
-    "for key, label in [(\"task1\",\"Task 1\"),(\"task2\",\"Task 2\"),(\"task3\",\"Task 3\"),(\"overall\",\"Overall\")]:\n",
-    "    b = baseline_scores.get(key, 0)\n",
-    "    a = trained_scores.get(key, 0)\n",
-    "    d = a - b\n",
-    "    print(f'{label:<10} {b:>8.3f} {a:>8.3f} {d:>+8.3f}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "\n",
-    "episodes   = [x[0] for x in training_log]\n",
-    "task_ids   = [x[1] for x in training_log]\n",
-    "ep_rewards = [x[2] for x in training_log]\n",
-    "\n",
-    "def moving_avg(data, window=5):\n",
-    "    return np.convolve(data, np.ones(window)/window, mode='valid')\n",
-    "\n",
-    "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n",
-    "fig.suptitle('Support Ticket Env - GRPO Training Results', fontsize=14, fontweight='bold')\n",
-    "\n",
-    "ax1 = axes[0]\n",
-    "colors = {1: '#3498db', 2: '#2ecc71', 3: '#e74c3c'}\n",
-    "for tid in [1, 2, 3]:\n",
-    "    mask = [i for i, t in enumerate(task_ids) if t == tid]\n",
-    "    if mask:\n",
-    "        x = [episodes[i] for i in mask]\n",
-    "        y = [ep_rewards[i] for i in mask]\n",
-    "        ax1.scatter(x, y, alpha=0.3, color=colors[tid], s=15)\n",
-    "        if len(y) >= 5:\n",
-    "            smoothed = moving_avg(y)\n",
-    "            ax1.plot(x[2:-2], smoothed, color=colors[tid], linewidth=2, label=f'Task {tid}')\n",
-    "        else:\n",
-    "            ax1.plot(x, y, color=colors[tid], linewidth=2, label=f'Task {tid}')\n",
-    "\n",
-    "ax1.set_xlabel('Episode')\n",
-    "ax1.set_ylabel('Avg Reward')\n",
-    "ax1.set_title('Training Reward per Episode')\n",
-    "ax1.legend()\n",
-    "ax1.grid(True, alpha=0.3)\n",
-    "ax1.set_ylim(-0.1, 1.1)\n",
-    "\n",
-    "ax2 = axes[1]\n",
-    "tasks       = ['Task 1', 'Task 2', 'Task 3', 'Overall']\n",
-    "keys        = ['task1', 'task2', 'task3', 'overall']\n",
-    "before_vals = [baseline_scores.get(k, 0) for k in keys]\n",
-    "after_vals  = [trained_scores.get(k, 0) for k in keys]\n",
-    "\n",
-    "x     = np.arange(len(tasks))\n",
-    "width = 0.35\n",
-    "\n",
-    "bars1 = ax2.bar(x - width/2, before_vals, width, label='Before Training', color='#95a5a6')\n",
-    "bars2 = ax2.bar(x + width/2, after_vals,  width, label='After GRPO',      color='#2ecc71')\n",
-    "\n",
-    "for bar in bars1:\n",
-    "    ax2.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.01,\n",
-    "             f'{bar.get_height():.2f}', ha='center', va='bottom', fontsize=9)\n",
-    "for bar in bars2:\n",
-    "    ax2.text(bar.get_x() + bar.get_width()/2., bar.get_height() + 0.01,\n",
-    "             f'{bar.get_height():.2f}', ha='center', va='bottom', fontsize=9,\n",
-    "             fontweight='bold', color='#27ae60')\n",
-    "\n",
-    "ax2.set_xticks(x)\n",
-    "ax2.set_xticklabels(tasks)\n",
-    "ax2.set_ylabel('Score (0-1)')\n",
-    "ax2.set_title('Before vs After GRPO Training')\n",
-    "ax2.legend()\n",
-    "ax2.grid(True, alpha=0.3, axis='y')\n",
-    "ax2.set_ylim(0, 1.15)\n",
-    "\n",
-    "plt.tight_layout()\n",
-    "plt.savefig('/content/grpo_results.png', dpi=150, bbox_inches='tight')\n",
-    "plt.show()\n",
-    "print('Chart saved to /content/grpo_results.png')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "os.makedirs(OUTPUT_DIR, exist_ok=True)\n",
-    "\n",
-    "model.save_pretrained(OUTPUT_DIR)\n",
-    "tokenizer.save_pretrained(OUTPUT_DIR)\n",
-    "print(f'Model saved to {OUTPUT_DIR}')\n",
-    "\n",
-    "try:\n",
-    "    from huggingface_hub import HfApi\n",
-    "    api = HfApi(token=HF_TOKEN)\n",
-    "    api.create_repo(HF_REPO_ID, exist_ok=True, private=False)\n",
-    "    api.upload_folder(folder_path=OUTPUT_DIR, repo_id=HF_REPO_ID, repo_type='model')\n",
-    "    api.upload_file(path_or_fileobj='/content/grpo_results.png', path_in_repo='grpo_results.png', repo_id=HF_REPO_ID, repo_type='model')\n",
-    "    print(f'Model pushed to: https://huggingface.co/{HF_REPO_ID}')\n",
-    "except Exception as e:\n",
-    "    print(f'Push failed: {e}')\n",
-    "    print(f'Model is saved locally at {OUTPUT_DIR}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from google.colab import files\n",
-    "files.download('/content/grpo_results.png')\n",
-    "\n",
-    "print('\\n' + '='*50)\n",
-    "print('FINAL TRAINING SUMMARY')\n",
-    "print('='*50)\n",
-    "print(f'Model:     {MODEL_NAME}')\n",
-    "print(f'Algorithm: GRPO')\n",
-    "print(f'Episodes:  {N_EPISODES}')\n",
-    "print(f'Env:       {ENV_BASE_URL}')\n",
-    "print()\n",
-    "print(f'{\"Task\":<10} {\"Before\":>8} {\"After\":>8} {\"Delta\":>8}')\n",
-    "print('-' * 38)\n",
-    "for key, label in [(\"task1\",\"Task 1\"),(\"task2\",\"Task 2\"),(\"task3\",\"Task 3\"),(\"overall\",\"Overall\")]:\n",
-    "    b = baseline_scores.get(key, 0)\n",
-    "    a = trained_scores.get(key, 0)\n",
-    "    d = a - b\n",
-    "    print(f'{label:<10} {b:>8.3f} {a:>8.3f} {d:>+8.3f}')\n",
-    "print('='*50)\n",
-    "print(f'Model: https://huggingface.co/{HF_REPO_ID}')"
-   ]
-  }
- ]
-}