updated grpo logic

README.md

@@ -1,14 +1,3 @@
----
-title: MetaGuard Ad Policy Sandbox
-emoji: 🛡
-colorFrom: blue
-colorTo: indigo
-sdk: docker
-app_port: 8000
-pinned: false
-license: mit
----
-
 # MetaGuard: A Multi-App RL Environment for Enterprise Ad Policy Compliance
 
 > An OpenEnv-compatible reinforcement learning environment that forces an LLM agent

grpo_train.ipynb

@@ -1,540 +1,653 @@
 {
- "nbformat": 4,
- "nbformat_minor": 0,
- "metadata": {
-  "colab": {
-   "provenance": [],
-   "gpuType": "A100"
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# 🛡️ MetaGuard — GRPO Training Notebook\n",
+        "\n",
+        "**Team:** Parth Singhal, Mehakveer Kaur, Kartik Goyal  \n",
+        "**HF Space:** https://huggingface.co/spaces/parth-1/MetaGuard  \n",
+        "**Hackathon:** OpenEnv — Meta × Scaler  \n",
+        "\n",
+        "This notebook trains **Llama 3.1 8B** using GRPO on the MetaGuard Ad Policy Compliance environment.\n",
+        "\n",
+        "### What this trains:\n",
+        "- Agent learns to follow structured SOP: `query_regulations → gather signals → submit_audit → decide`\n",
+        "- Reward shaped by correctness, sequence compliance, API failure recovery\n",
+        "- Environment runs locally in the notebook (fast); GPU handles only the model"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 1 — Install Dependencies"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "!pip install unsloth trl transformers datasets accelerate peft -q\n",
+        "!pip install openenv-core==0.2.1 --no-deps -q\n",
+        "!pip install fastapi uvicorn pydantic requests openai matplotlib -q\n",
+        "print('✅ Dependencies installed')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 2 — Clone Repo"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import os\n",
+        "\n",
+        "if not os.path.exists('meta-ad-policy-sandbox'):\n",
+        "    !git clone https://github.com/Parth380/meta-ad-policy-sandbox.git\n",
+        "\n",
+        "%cd meta-ad-policy-sandbox\n",
+        "!pip install -e . -q\n",
+        "os.makedirs('outputs', exist_ok=True)\n",
+        "print('Repo installed & outputs/ ready')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 3 — Config (SET THESE)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import os\n",
+        "\n",
+        "os.environ['ENV_URL']  = 'http://localhost:8000'  # local env (fast); change to HF Space URL if needed\n",
+        "os.environ['HF_REPO']  = 'parth-1/metaguard-llama3.1-8b-grpo'\n",
+        "os.environ['HF_TOKEN'] = ''  # paste your HF write token here\n",
+        "\n",
+        "ENV_URL   = os.environ['ENV_URL']\n",
+        "HF_TOKEN  = os.environ['HF_TOKEN']\n",
+        "HF_REPO   = os.environ['HF_REPO']\n",
+        "\n",
+        "print(f'ENV_URL  : {ENV_URL}')\n",
+        "print(f'HF_REPO  : {HF_REPO}')\n",
+        "print(f'HF_TOKEN : {\"set\" if HF_TOKEN else \"MISSING -- set above before Cell 11\"}')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 4 — Boot Local Environment"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import subprocess, time, threading, requests\n",
+        "import uvicorn\n",
+        "\n",
+        "procs = [\n",
+        "    subprocess.Popen(['python', 'apps/regulatory_api.py']),\n",
+        "    subprocess.Popen(['python', 'apps/crm_api.py']),\n",
+        "    subprocess.Popen(['python', 'apps/audit_api.py']),\n",
+        "]\n",
+        "time.sleep(3)\n",
+        "\n",
+        "from server.app import app as _env_app\n",
+        "threading.Thread(\n",
+        "    target=uvicorn.run,\n",
+        "    kwargs={'app': _env_app, 'host': '0.0.0.0', 'port': 8000, 'log_level': 'warning'},\n",
+        "    daemon=True,\n",
+        ").start()\n",
+        "time.sleep(2)\n",
+        "\n",
+        "for i in range(20):\n",
+        "    try:\n",
+        "        r = requests.post(f'{ENV_URL}/reset', json={'task_id': 'task_1_healthcare'}, timeout=5)\n",
+        "        if r.status_code == 200:\n",
+        "            print(f'Environment ready (attempt {i+1})')\n",
+        "            break\n",
+        "    except:\n",
+        "        pass\n",
+        "    time.sleep(1)\n",
+        "else:\n",
+        "    raise RuntimeError('ENV not reachable after 20 attempts')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 5 — Imports + Helpers"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import json\n",
+        "import random\n",
+        "import torch\n",
+        "import matplotlib.pyplot as plt\n",
+        "from collections import defaultdict\n",
+        "\n",
+        "from datasets import Dataset\n",
+        "from unsloth import FastLanguageModel, PatchFastRL\n",
+        "from trl import GRPOTrainer, GRPOConfig\n",
+        "\n",
+        "PatchFastRL('GRPO', FastLanguageModel)\n",
+        "\n",
+        "ALLOWED_ACTIONS = [\n",
+        "    'query_regulations', 'analyze_image', 'check_advertiser_history',\n",
+        "    'request_landing_page', 'request_id_verification',\n",
+        "    'submit_audit', 'approve', 'reject',\n",
+        "]\n",
+        "\n",
+        "class EnvClient:\n",
+        "    def __init__(self, url):\n",
+        "        self.url = url\n",
+        "    def reset(self, task_id):\n",
+        "        return requests.post(f'{self.url}/reset', json={'task_id': task_id}, timeout=8).json()\n",
+        "    def step(self, action):\n",
+        "        return requests.post(f'{self.url}/step', json={'action': action}, timeout=8).json()\n",
+        "\n",
+        "def safe_step(client, action):\n",
+        "    for _ in range(3):\n",
+        "        try:\n",
+        "            return client.step(action)\n",
+        "        except:\n",
+        "            time.sleep(0.5)\n",
+        "    return {'reward': -0.3}\n",
+        "\n",
+        "def extract_json(text):\n",
+        "    try:\n",
+        "        if '```' in text:\n",
+        "            text = text.split('```')[1]\n",
+        "            if text.startswith('json'):\n",
+        "                text = text[4:]\n",
+        "        return json.loads(text.strip())\n",
+        "    except:\n",
+        "        return None\n",
+        "\n",
+        "print('✅ Helpers loaded')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 6 — Dataset"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "BASE_SCENARIOS = [\n",
+        "    {'task_id': 'task_1_healthcare',\n",
+        "     'text': \"Healthcare ad: 'miracle supplement cures disease'. No actions taken yet.\",\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_2_financial',\n",
+        "     'text': \"Financial ad: 'guaranteed 500% returns, zero risk'. No actions taken yet.\",\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_3_multimodal',\n",
+        "     'text': 'Multimodal ad: image may contain hidden violation. No actions taken yet.',\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_6_conflict',\n",
+        "     'text': 'High-trust advertiser but policy borderline. No actions taken yet.',\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_7_ambiguous',\n",
+        "     'text': 'Ambiguous wellness ad. Policy confidence low. No actions taken yet.',\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_8_adversarial',\n",
+        "     'text': 'Natural supplement ad. Image may hide violation. No actions taken yet.',\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_9_dependency_trap',\n",
+        "     'text': 'Certified wellness product. Regulations and CRM look clean. No actions taken yet.',\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    {'task_id': 'task_10_failure',\n",
+        "     'text': 'Healthcare ad. First API call may fail. No actions taken yet.',\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    # task_4 targeting — fresh\n",
+        "    {'task_id': 'task_4_targeting',\n",
+        "     'text': \"Financial ad targeting young users: 'Start Your First Investment Portfolio'. No actions taken yet.\",\n",
+        "     'actions_already_taken': [], 'setup_actions': []},\n",
+        "    # task_4 targeting — mid state\n",
+        "    {'task_id': 'task_4_targeting',\n",
+        "     'text': 'Financial ad targeting young users. Policy queried, need to verify age targeting.',\n",
+        "     'actions_already_taken': ['query_regulations'],\n",
+        "     'setup_actions': [{'action_type': 'query_regulations', 'reasoning': 'policy lookup'}]},\n",
+        "    # task_4 targeting — audit ready\n",
+        "    {'task_id': 'task_4_targeting',\n",
+        "     'text': 'Financial ad targeting minors. Policy, history, and ID verification done. Submit audit.',\n",
+        "     'actions_already_taken': ['query_regulations', 'check_advertiser_history', 'request_id_verification'],\n",
+        "     'setup_actions': [\n",
+        "         {'action_type': 'query_regulations', 'reasoning': 'policy lookup'},\n",
+        "         {'action_type': 'check_advertiser_history', 'reasoning': 'trust score'},\n",
+        "         {'action_type': 'request_id_verification', 'reasoning': 'age check'},\n",
+        "     ]},\n",
+        "    {'task_id': 'task_1_healthcare',\n",
+        "     'text': 'Healthcare ad. Policy already queried.',\n",
+        "     'actions_already_taken': ['query_regulations'],\n",
+        "     'setup_actions': [{'action_type': 'query_regulations', 'reasoning': 'policy lookup'}]},\n",
+        "    {'task_id': 'task_2_financial',\n",
+        "     'text': 'Financial ad. Policy and history checked. Submit audit next.',\n",
+        "     'actions_already_taken': ['query_regulations', 'check_advertiser_history'],\n",
+        "     'setup_actions': [\n",
+        "         {'action_type': 'query_regulations', 'reasoning': 'policy lookup'},\n",
+        "         {'action_type': 'check_advertiser_history', 'reasoning': 'trust score'},\n",
+        "     ]},\n",
+        "    {'task_id': 'task_2_financial',\n",
+        "     'text': 'Financial ad. Policy, history, audit all complete. Make final decision.',\n",
+        "     'actions_already_taken': ['query_regulations', 'check_advertiser_history', 'submit_audit'],\n",
+        "     'setup_actions': [\n",
+        "         {'action_type': 'query_regulations', 'reasoning': 'policy lookup'},\n",
+        "         {'action_type': 'check_advertiser_history', 'reasoning': 'trust score'},\n",
+        "         {'action_type': 'submit_audit', 'reasoning': 'audit log'},\n",
+        "     ]},\n",
+        "]\n",
+        "\n",
+        "PROMPT_TEMPLATE = \"\"\"You are an enterprise Ad Policy Compliance Agent.\n",
+        "\n",
+        "You MUST choose exactly ONE action_type from this list (any other value is invalid):\n",
+        "- query_regulations\n",
+        "- analyze_image\n",
+        "- check_advertiser_history\n",
+        "- request_landing_page\n",
+        "- request_id_verification\n",
+        "- submit_audit\n",
+        "- approve\n",
+        "- reject\n",
+        "\n",
+        "REQUIRED PHASE ORDER:\n",
+        "1. query_regulations  -> always first\n",
+        "2. analyze_image / check_advertiser_history  -> gather signals\n",
+        "3. submit_audit  -> always before final decision\n",
+        "4. approve OR reject  -> only after audit\n",
+        "\n",
+        "HARD RULES:\n",
+        "- NEVER repeat an action listed in `actions_already_taken`.\n",
+        "- Respond with ONLY a valid JSON object. No markdown, no prose.\n",
+        "\n",
+        "Required format:\n",
+        "{{\\\"action_type\\\": \\\"<one_of_the_actions_above>\\\", \\\"reasoning\\\": \\\"<short reason>\\\"}}\n",
+        "\n",
+        "Scenario: {text}\n",
+        "actions_already_taken: {actions_already_taken}\n",
+        "\n",
+        "Your next action?\"\"\"\n",
+        "\n",
+        "def build_dataset():\n",
+        "    rows = []\n",
+        "    for s in BASE_SCENARIOS:\n",
+        "        prompt = PROMPT_TEMPLATE.format(\n",
+        "            text=s['text'],\n",
+        "            actions_already_taken=json.dumps(s['actions_already_taken']),\n",
+        "        )\n",
+        "        rows.append({'prompt': prompt, 'task_id': s['task_id'], 'setup_actions': s['setup_actions']})\n",
+        "    return Dataset.from_list(rows * 10)\n",
+        "\n",
+        "dataset = build_dataset()\n",
+        "print(f'✅ Dataset: {len(dataset)} examples')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 7 — Reward Function"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "# Track rewards for plotting\n",
+        "reward_log = []\n",
+        "step_log = []\n",
+        "global_step_counter = [0]\n",
+        "\n",
+        "def reward_environment(prompts, completions, task_id=None, setup_actions=None, **kwargs):\n",
+        "    client = EnvClient(ENV_URL)\n",
+        "    rewards = []\n",
+        "\n",
+        "    for completion, t_id, setup in zip(completions, task_id, setup_actions):\n",
+        "        parsed = extract_json(completion)\n",
+        "        if not parsed:\n",
+        "            rewards.append(-1.0)\n",
+        "            continue\n",
+        "\n",
+        "        action_type = parsed.get('action_type')\n",
+        "        if action_type not in ALLOWED_ACTIONS:\n",
+        "            rewards.append(-1.0)\n",
+        "            continue\n",
+        "\n",
+        "        action = {\n",
+        "            'action_type': action_type,\n",
+        "            'reasoning': parsed.get('reasoning', 'format-compliant'),\n",
+        "        }\n",
+        "\n",
+        "        try:\n",
+        "            client.reset(t_id)\n",
+        "            for s in setup:\n",
+        "                safe_step(client, s)\n",
+        "\n",
+        "            result = safe_step(client, action)\n",
+        "            env_reward = float(result.get('reward', -0.2))\n",
+        "            status_msg = (result.get('status_message') or '').lower()\n",
+        "\n",
+        "            rejected = (\n",
+        "                'api failure' in status_msg\n",
+        "                or 'invalid action' in status_msg\n",
+        "                or 'must call' in status_msg\n",
+        "            )\n",
+        "            shaped = -0.5 if rejected else 0.5 + env_reward\n",
+        "            rewards.append(shaped)\n",
+        "\n",
+        "        except Exception:\n",
+        "            rewards.append(-0.3)\n",
+        "\n",
+        "    # Log for plot\n",
+        "    avg = sum(rewards) / len(rewards) if rewards else 0.0\n",
+        "    global_step_counter[0] += 1\n",
+        "    reward_log.append(avg)\n",
+        "    step_log.append(global_step_counter[0])\n",
+        "\n",
+        "    return rewards\n",
+        "\n",
+        "print('✅ Reward function ready')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 8 — Load Model"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "if torch.cuda.is_available():\n",
+        "    _props = torch.cuda.get_device_properties(0)\n",
+        "    _vram = _props.total_memory\n",
+        "    _cc = (_props.major, _props.minor)\n",
+        "    print(f'GPU: {_props.name}  VRAM: {_vram / 1024**3:.1f} GB  Compute: {_cc[0]}.{_cc[1]}')\n",
+        "else:\n",
+        "    _vram, _cc = 0, (0, 0)\n",
+        "\n",
+        "USE_4BIT = _vram < 40 * 1024**3   # T4/L4 → 4-bit; A100 → full precision\n",
+        "USE_BF16 = _cc >= (8, 0) and not USE_4BIT  # bf16 only with full-precision weights; 4-bit LoRA uses fp16\n",
+        "print(f'4-bit: {USE_4BIT}  bf16: {USE_BF16}')\n",
+        "\n",
+        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name='unsloth/Llama-3.1-8B-Instruct',\n",
+        "    load_in_4bit=USE_4BIT,\n",
+        "    max_seq_length=2048,\n",
+        "    dtype=torch.float16 if USE_4BIT else None,\n",
+        ")\n",
+        "\n",
+        "model = FastLanguageModel.get_peft_model(\n",
+        "    model,\n",
+        "    r=32,\n",
+        "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
+        "    lora_alpha=64,\n",
+        "    lora_dropout=0,\n",
+        "    bias='none',\n",
+        "    use_gradient_checkpointing='unsloth',\n",
+        "    random_state=3407,\n",
+        ")\n",
+        "print('✅ Model loaded')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 9 — Train"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "trainer = GRPOTrainer(\n",
+        "    model=model,\n",
+        "    reward_funcs=[reward_environment],\n",
+        "    args=GRPOConfig(\n",
+        "        output_dir='outputs',\n",
+        "        learning_rate=2e-5,\n",
+        "        num_train_epochs=3,\n",
+        "        per_device_train_batch_size=2 if not USE_4BIT else 1,\n",
+        "        gradient_accumulation_steps=4,\n",
+        "        num_generations=4 if not USE_4BIT else 2,\n",
+        "        max_prompt_length=768,\n",
+        "        max_completion_length=128,\n",
+        "        logging_steps=5,\n",
+        "        warmup_steps=10,\n",
+        "        bf16=USE_BF16,\n",
+        "        fp16=not USE_BF16,\n",
+        "        report_to='none',\n",
+        "    ),\n",
+        "    train_dataset=dataset,\n",
+        "    tokenizer=tokenizer,\n",
+        ")\n",
+        "\n",
+        "print('Starting GRPO training...')\n",
+        "print(f'  bf16={USE_BF16}  fp16={not USE_BF16}  batch={2 if not USE_4BIT else 1}  gens={4 if not USE_4BIT else 2}')\n",
+        "trainer.train()\n",
+        "print('Training complete')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 10 — Plot Reward Curve"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "import numpy as np\n",
+        "import pandas as pd\n",
+        "import os\n",
+        "\n",
+        "os.makedirs('outputs', exist_ok=True)\n",
+        "\n",
+        "def moving_avg(data, window=5):\n",
+        "    if len(data) < window:\n",
+        "        return data\n",
+        "    return list(np.convolve(data, np.ones(window)/window, mode='valid'))\n",
+        "\n",
+        "hist = pd.DataFrame(trainer.state.log_history)\n",
+        "\n",
+        "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n",
+        "\n",
+        "# --- Plot 1: Reward curve (from our custom log) ---\n",
+        "ax = axes[0]\n",
+        "ax.plot(step_log, reward_log, alpha=0.3, color='steelblue', label='Raw')\n",
+        "smoothed = moving_avg(reward_log)\n",
+        "ax.plot(range(len(smoothed)), smoothed, color='steelblue', linewidth=2, label='Smoothed (MA-5)')\n",
+        "ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.8)\n",
+        "ax.set_xlabel('Reward Eval Step')\n",
+        "ax.set_ylabel('Avg Reward per Batch')\n",
+        "ax.set_title('Reward Curve')\n",
+        "ax.legend()\n",
+        "ax.grid(alpha=0.3)\n",
+        "\n",
+        "# --- Plot 2: Loss curve (from trainer logs) ---\n",
+        "ax = axes[1]\n",
+        "loss_rows = hist.dropna(subset=['loss']) if 'loss' in hist.columns else pd.DataFrame()\n",
+        "if not loss_rows.empty:\n",
+        "    ax.plot(loss_rows['step'], loss_rows['loss'], color='#7c3aed', linewidth=2)\n",
+        "    ax.set_xlabel('Training Step')\n",
+        "    ax.set_ylabel('Loss')\n",
+        "    ax.set_title('GRPO Loss')\n",
+        "    ax.grid(alpha=0.3)\n",
+        "else:\n",
+        "    ax.text(0.5, 0.5, 'No loss data logged', ha='center', va='center', transform=ax.transAxes)\n",
+        "    ax.set_title('GRPO Loss')\n",
+        "\n",
+        "# --- Plot 3: Reward from trainer logs (if available) ---\n",
+        "ax = axes[2]\n",
+        "reward_cols = [c for c in hist.columns if 'reward' in c.lower() and 'std' not in c.lower()]\n",
+        "if reward_cols:\n",
+        "    col = reward_cols[0]\n",
+        "    rr = hist.dropna(subset=[col])\n",
+        "    ax.plot(rr['step'], rr[col], color='#16a34a', linewidth=2)\n",
+        "    ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.8)\n",
+        "    ax.set_xlabel('Training Step')\n",
+        "    ax.set_ylabel(col)\n",
+        "    ax.set_title('Trainer Reward Log')\n",
+        "    ax.grid(alpha=0.3)\n",
+        "else:\n",
+        "    ax.text(0.5, 0.5, 'No trainer reward data', ha='center', va='center', transform=ax.transAxes)\n",
+        "    ax.set_title('Trainer Reward Log')\n",
+        "\n",
+        "plt.tight_layout()\n",
+        "plt.savefig('outputs/training_plots.png', dpi=150)\n",
+        "plt.show()\n",
+        "print('Saved to outputs/training_plots.png')\n",
+        "\n",
+        "n = len(reward_log)\n",
+        "first_10 = reward_log[:min(10, n)]\n",
+        "last_10  = reward_log[max(0, n-10):]\n",
+        "print(f'\\n--- Before vs After ---')\n",
+        "print(f'Avg reward (first 10 steps): {sum(first_10)/len(first_10):.3f}')\n",
+        "print(f'Avg reward (last 10 steps) : {sum(last_10)/len(last_10):.3f}')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 11 — Before vs After: Baseline Comparison"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "from unsloth import FastLanguageModel as FLM\n",
+        "\n",
+        "FLM.for_inference(model)\n",
+        "\n",
+        "test_scenarios = [\n",
+        "    ('task_1_healthcare', \"Healthcare ad: 'miracle cure'. No actions taken yet.\", []),\n",
+        "    ('task_2_financial', \"Financial ad: 'guaranteed returns'. No actions taken yet.\", []),\n",
+        "    ('task_4_targeting', \"Financial ad targeting teens. No actions taken yet.\", []),\n",
+        "    ('task_2_financial', \"Financial ad. Policy, history, audit done. Decide.\",\n",
+        "     ['query_regulations', 'check_advertiser_history', 'submit_audit']),\n",
+        "]\n",
+        "\n",
+        "print('=== Trained Model Outputs ===\\n')\n",
+        "for task, text, taken in test_scenarios:\n",
+        "    prompt = PROMPT_TEMPLATE.format(text=text, actions_already_taken=json.dumps(taken))\n",
+        "    inputs = tokenizer(prompt, return_tensors='pt').to('cuda')\n",
+        "    out = model.generate(**inputs, max_new_tokens=64, temperature=0.1, do_sample=True)\n",
+        "    decoded = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
+        "    parsed = extract_json(decoded) or decoded.strip()[:120]\n",
+        "    print(f'[{task}] taken={taken}')\n",
+        "    print(f'  -> {parsed}\\n')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Cell 12 — Save + Push to HF Hub"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": [
+        "model.save_pretrained('outputs/lora_adapter')\n",
+        "tokenizer.save_pretrained('outputs/lora_adapter')\n",
+        "print('LoRA adapter saved')\n",
+        "\n",
+        "print('Merging adapter into base model...')\n",
+        "merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name='outputs/lora_adapter',\n",
+        "    load_in_4bit=False,\n",
+        "    max_seq_length=2048,\n",
+        ")\n",
+        "merged_model.save_pretrained_merged(\n",
+        "    'outputs/merged',\n",
+        "    merged_tokenizer,\n",
+        "    save_method='merged_16bit',\n",
+        ")\n",
+        "print('Merged model saved to outputs/merged')\n",
+        "\n",
+        "if HF_REPO and HF_TOKEN:\n",
+        "    print(f'Pushing to {HF_REPO}...')\n",
+        "    merged_model.push_to_hub_merged(\n",
+        "        HF_REPO,\n",
+        "        merged_tokenizer,\n",
+        "        save_method='merged_16bit',\n",
+        "        token=HF_TOKEN,\n",
+        "    )\n",
+        "    print(f'Model live at https://huggingface.co/{HF_REPO}')\n",
+        "else:\n",
+        "    print('Set HF_REPO and HF_TOKEN in Cell 3 to push to Hub')\n",
+        "\n",
+        "print('Done.')"
+      ],
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "A100"
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
   },
-  "kernelspec": {
-   "display_name": "Python 3",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python"
-  },
-  "accelerator": "GPU"
- },
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# 🛡️ MetaGuard — GRPO Training Notebook\n",
-    "\n",
-    "**Team:** Parth Singhal, Mehakveer Kaur, Kartik Goyal  \n",
-    "**HF Space:** https://huggingface.co/spaces/parth-1/MetaGuard  \n",
-    "**Hackathon:** OpenEnv — Meta × Scaler  \n",
-    "\n",
-    "This notebook trains **Llama 3.1 8B** using GRPO on the MetaGuard Ad Policy Compliance environment.\n",
-    "\n",
-    "### What this trains:\n",
-    "- Agent learns to follow structured SOP: `query_regulations → gather signals → submit_audit → decide`\n",
-    "- Reward shaped by correctness, sequence compliance, API failure recovery\n",
-    "- Environment hosted on HF Space (A100 runs both training + env)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 1 — Install Dependencies"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "!pip install unsloth trl transformers datasets accelerate peft -q\n",
-    "!pip install openenv-core==0.2.1 --no-deps -q\n",
-    "!pip install fastapi uvicorn pydantic requests openai matplotlib -q\n",
-    "print('✅ Dependencies installed')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 2 — Clone Repo"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "if not os.path.exists('meta-ad-policy-sandbox'):\n",
-    "    !git clone https://github.com/Parth380/meta-ad-policy-sandbox.git\n",
-    "\n",
-    "%cd meta-ad-policy-sandbox\n",
-    "print('✅ Repo ready')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 3 — Config (SET THESE)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "\n",
-    "os.environ['ENV_URL']  = 'https://parth-1-metaguard.hf.space'  # your HF Space URL\n",
-    "os.environ['HF_REPO']  = 'parth-1/metaguard-llama3.1-8b-grpo'  # model push destination\n",
-    "os.environ['HF_TOKEN'] = ''  # ← paste your HF write token here\n",
-    "\n",
-    "ENV_URL   = os.environ['ENV_URL']\n",
-    "HF_TOKEN  = os.environ['HF_TOKEN']\n",
-    "HF_REPO   = os.environ['HF_REPO']\n",
-    "\n",
-    "print(f'ENV_URL  : {ENV_URL}')\n",
-    "print(f'HF_REPO  : {HF_REPO}')\n",
-    "print(f'HF_TOKEN : {\"set ✅\" if HF_TOKEN else \"MISSING ❌\"}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 4 — Wake Up HF Space"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import requests, time\n",
-    "\n",
-    "print('Waking up HF Space...')\n",
-    "for i in range(20):\n",
-    "    try:\n",
-    "        r = requests.post(\n",
-    "            f\"{ENV_URL}/reset\",\n",
-    "            json={'task_id': 'task_1_healthcare'},\n",
-    "            timeout=10\n",
-    "        )\n",
-    "        if r.status_code == 200:\n",
-    "            print(f'✅ Environment ready (attempt {i+1})')\n",
-    "            break\n",
-    "    except Exception as e:\n",
-    "        print(f'  attempt {i+1}: waiting... ({e})')\n",
-    "    time.sleep(3)\n",
-    "else:\n",
-    "    raise RuntimeError('❌ ENV not reachable after 20 attempts')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 5 — Imports + Helpers"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import json\n",
-    "import random\n",
-    "import torch\n",
-    "import matplotlib.pyplot as plt\n",
-    "from collections import defaultdict\n",
-    "\n",
-    "from datasets import Dataset\n",
-    "from unsloth import FastLanguageModel, PatchFastRL\n",
-    "from trl import GRPOTrainer, GRPOConfig\n",
-    "\n",
-    "PatchFastRL('GRPO', FastLanguageModel)\n",
-    "\n",
-    "ALLOWED_ACTIONS = [\n",
-    "    'query_regulations', 'analyze_image', 'check_advertiser_history',\n",
-    "    'request_landing_page', 'request_id_verification',\n",
-    "    'submit_audit', 'approve', 'reject',\n",
-    "]\n",
-    "\n",
-    "class EnvClient:\n",
-    "    def __init__(self, url):\n",
-    "        self.url = url\n",
-    "    def reset(self, task_id):\n",
-    "        return requests.post(f'{self.url}/reset', json={'task_id': task_id}, timeout=8).json()\n",
-    "    def step(self, action):\n",
-    "        return requests.post(f'{self.url}/step', json={'action': action}, timeout=8).json()\n",
-    "\n",
-    "def safe_step(client, action):\n",
-    "    for _ in range(3):\n",
-    "        try:\n",
-    "            return client.step(action)\n",
-    "        except:\n",
-    "            time.sleep(0.5)\n",
-    "    return {'reward': -0.3}\n",
-    "\n",
-    "def extract_json(text):\n",
-    "    try:\n",
-    "        if '```' in text:\n",
-    "            text = text.split('```')[1]\n",
-    "            if text.startswith('json'):\n",
-    "                text = text[4:]\n",
-    "        return json.loads(text.strip())\n",
-    "    except:\n",
-    "        return None\n",
-    "\n",
-    "print('✅ Helpers loaded')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 6 — Dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "BASE_SCENARIOS = [\n",
-    "    {'task_id': 'task_1_healthcare',\n",
-    "     'text': \"Healthcare ad: 'miracle supplement cures disease'. No actions taken yet.\",\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_2_financial',\n",
-    "     'text': \"Financial ad: 'guaranteed 500% returns, zero risk'. No actions taken yet.\",\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_3_multimodal',\n",
-    "     'text': 'Multimodal ad: image may contain hidden violation. No actions taken yet.',\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_6_conflict',\n",
-    "     'text': 'High-trust advertiser but policy borderline. No actions taken yet.',\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_7_ambiguous',\n",
-    "     'text': 'Ambiguous wellness ad. Policy confidence low. No actions taken yet.',\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_8_adversarial',\n",
-    "     'text': 'Natural supplement ad. Image may hide violation. No actions taken yet.',\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_9_dependency_trap',\n",
-    "     'text': 'Certified wellness product. Regulations and CRM look clean. No actions taken yet.',\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_10_failure',\n",
-    "     'text': 'Healthcare ad. First API call may fail. No actions taken yet.',\n",
-    "     'actions_already_taken': [], 'setup_actions': []},\n",
-    "    {'task_id': 'task_1_healthcare',\n",
-    "     'text': 'Healthcare ad. Policy already queried.',\n",
-    "     'actions_already_taken': ['query_regulations'],\n",
-    "     'setup_actions': [{'action_type': 'query_regulations', 'reasoning': 'policy lookup'}]},\n",
-    "    {'task_id': 'task_2_financial',\n",
-    "     'text': 'Financial ad. Policy and history checked. Submit audit next.',\n",
-    "     'actions_already_taken': ['query_regulations', 'check_advertiser_history'],\n",
-    "     'setup_actions': [\n",
-    "         {'action_type': 'query_regulations', 'reasoning': 'policy lookup'},\n",
-    "         {'action_type': 'check_advertiser_history', 'reasoning': 'trust score'},\n",
-    "     ]},\n",
-    "    {'task_id': 'task_2_financial',\n",
-    "     'text': 'Financial ad. Policy, history, audit all complete. Make final decision.',\n",
-    "     'actions_already_taken': ['query_regulations', 'check_advertiser_history', 'submit_audit'],\n",
-    "     'setup_actions': [\n",
-    "         {'action_type': 'query_regulations', 'reasoning': 'policy lookup'},\n",
-    "         {'action_type': 'check_advertiser_history', 'reasoning': 'trust score'},\n",
-    "         {'action_type': 'submit_audit', 'reasoning': 'audit log'},\n",
-    "     ]},\n",
-    "]\n",
-    "\n",
-    "PROMPT_TEMPLATE = \"\"\"You are an enterprise Ad Policy Compliance Agent.\n",
-    "\n",
-    "You MUST choose exactly ONE action_type from this list (any other value is invalid):\n",
-    "- query_regulations\n",
-    "- analyze_image\n",
-    "- check_advertiser_history\n",
-    "- request_landing_page\n",
-    "- request_id_verification\n",
-    "- submit_audit\n",
-    "- approve\n",
-    "- reject\n",
-    "\n",
-    "REQUIRED PHASE ORDER:\n",
-    "1. query_regulations  -> always first\n",
-    "2. analyze_image / check_advertiser_history  -> gather signals\n",
-    "3. submit_audit  -> always before final decision\n",
-    "4. approve OR reject  -> only after audit\n",
-    "\n",
-    "HARD RULES:\n",
-    "- NEVER repeat an action listed in `actions_already_taken`.\n",
-    "- Respond with ONLY a valid JSON object. No markdown, no prose.\n",
-    "\n",
-    "Required format:\n",
-    "{{\\\"action_type\\\": \\\"<one_of_the_actions_above>\\\", \\\"reasoning\\\": \\\"<short reason>\\\"}}\n",
-    "\n",
-    "Scenario: {text}\n",
-    "actions_already_taken: {actions_already_taken}\n",
-    "\n",
-    "Your next action?\"\"\"\n",
-    "\n",
-    "def build_dataset():\n",
-    "    rows = []\n",
-    "    for s in BASE_SCENARIOS:\n",
-    "        prompt = PROMPT_TEMPLATE.format(\n",
-    "            text=s['text'],\n",
-    "            actions_already_taken=json.dumps(s['actions_already_taken']),\n",
-    "        )\n",
-    "        rows.append({'prompt': prompt, 'task_id': s['task_id'], 'setup_actions': s['setup_actions']})\n",
-    "    return Dataset.from_list(rows * 10)\n",
-    "\n",
-    "dataset = build_dataset()\n",
-    "print(f'✅ Dataset: {len(dataset)} examples')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 7 — Reward Function"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Track rewards for plotting\n",
-    "reward_log = []\n",
-    "step_log = []\n",
-    "global_step_counter = [0]\n",
-    "\n",
-    "def reward_environment(prompts, completions, task_id=None, setup_actions=None, **kwargs):\n",
-    "    client = EnvClient(ENV_URL)\n",
-    "    rewards = []\n",
-    "\n",
-    "    for completion, t_id, setup in zip(completions, task_id, setup_actions):\n",
-    "        parsed = extract_json(completion)\n",
-    "        if not parsed:\n",
-    "            rewards.append(-1.0)\n",
-    "            continue\n",
-    "\n",
-    "        action_type = parsed.get('action_type')\n",
-    "        if action_type not in ALLOWED_ACTIONS:\n",
-    "            rewards.append(-1.0)\n",
-    "            continue\n",
-    "\n",
-    "        action = {\n",
-    "            'action_type': action_type,\n",
-    "            'reasoning': parsed.get('reasoning', 'format-compliant'),\n",
-    "        }\n",
-    "\n",
-    "        try:\n",
-    "            client.reset(t_id)\n",
-    "            for s in setup:\n",
-    "                safe_step(client, s)\n",
-    "\n",
-    "            result = safe_step(client, action)\n",
-    "            env_reward = float(result.get('reward', -0.2))\n",
-    "            status_msg = (result.get('status_message') or '').lower()\n",
-    "\n",
-    "            rejected = (\n",
-    "                'api failure' in status_msg\n",
-    "                or 'invalid action' in status_msg\n",
-    "                or 'must call' in status_msg\n",
-    "            )\n",
-    "            shaped = -0.5 if rejected else 0.5 + env_reward\n",
-    "            rewards.append(shaped)\n",
-    "\n",
-    "        except Exception:\n",
-    "            rewards.append(-0.3)\n",
-    "\n",
-    "    # Log for plot\n",
-    "    avg = sum(rewards) / len(rewards) if rewards else 0.0\n",
-    "    global_step_counter[0] += 1\n",
-    "    reward_log.append(avg)\n",
-    "    step_log.append(global_step_counter[0])\n",
-    "\n",
-    "    return rewards\n",
-    "\n",
-    "print('✅ Reward function ready')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 8 — Load Model"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "USE_4BIT = not torch.cuda.is_available() or torch.cuda.get_device_properties(0).total_memory < 40 * 1024**3\n",
-    "print(f'GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')\n",
-    "print(f'4-bit quant: {USE_4BIT}')\n",
-    "\n",
-    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name='unsloth/Llama-3.1-8B-Instruct',\n",
-    "    load_in_4bit=USE_4BIT,\n",
-    "    max_seq_length=2048,\n",
-    "    dtype=None,\n",
-    ")\n",
-    "\n",
-    "model = FastLanguageModel.get_peft_model(\n",
-    "    model,\n",
-    "    r=32,\n",
-    "    target_modules=['q_proj', 'k_proj', 'v_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj'],\n",
-    "    lora_alpha=64,\n",
-    "    lora_dropout=0,\n",
-    "    bias='none',\n",
-    "    use_gradient_checkpointing='unsloth',\n",
-    "    random_state=3407,\n",
-    ")\n",
-    "print('✅ Model loaded')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 9 — Train"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "trainer = GRPOTrainer(\n",
-    "    model=model,\n",
-    "    reward_funcs=[reward_environment],\n",
-    "    args=GRPOConfig(\n",
-    "        output_dir='outputs',\n",
-    "        learning_rate=2e-5,\n",
-    "        num_train_epochs=3,\n",
-    "        per_device_train_batch_size=2,\n",
-    "        gradient_accumulation_steps=4,\n",
-    "        num_generations=4,\n",
-    "        max_prompt_length=768,\n",
-    "        max_completion_length=128,\n",
-    "        logging_steps=5,\n",
-    "        warmup_ratio=0.1,\n",
-    "        bf16=True,\n",
-    "        report_to='none',\n",
-    "    ),\n",
-    "    train_dataset=dataset,\n",
-    "    tokenizer=tokenizer,\n",
-    ")\n",
-    "\n",
-    "print('🚀 Starting GRPO training...')\n",
-    "trainer.train()\n",
-    "print('✅ Training complete')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 10 — Plot Reward Curve"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "\n",
-    "# Smooth with moving average\n",
-    "def moving_avg(data, window=5):\n",
-    "    if len(data) < window:\n",
-    "        return data\n",
-    "    return np.convolve(data, np.ones(window)/window, mode='valid')\n",
-    "\n",
-    "fig, ax = plt.subplots(figsize=(10, 5))\n",
-    "ax.plot(step_log, reward_log, alpha=0.3, color='steelblue', label='Raw reward')\n",
-    "smoothed = moving_avg(reward_log)\n",
-    "ax.plot(range(len(smoothed)), smoothed, color='steelblue', linewidth=2, label='Smoothed (MA-5)')\n",
-    "ax.axhline(y=0, color='gray', linestyle='--', linewidth=0.8)\n",
-    "ax.set_xlabel('Training Step')\n",
-    "ax.set_ylabel('Avg Reward per Batch')\n",
-    "ax.set_title('MetaGuard GRPO — Reward Curve')\n",
-    "ax.legend()\n",
-    "ax.grid(alpha=0.3)\n",
-    "\n",
-    "plt.tight_layout()\n",
-    "plt.savefig('outputs/reward_plot.png', dpi=150)\n",
-    "plt.show()\n",
-    "print('✅ Plot saved to outputs/reward_plot.png')\n",
-    "\n",
-    "# Print before/after summary\n",
-    "n = len(reward_log)\n",
-    "first_10 = reward_log[:min(10, n)]\n",
-    "last_10  = reward_log[max(0, n-10):]\n",
-    "print(f'\\n--- Results ---')\n",
-    "print(f'Avg reward (first 10 steps): {sum(first_10)/len(first_10):.3f}')\n",
-    "print(f'Avg reward (last 10 steps) : {sum(last_10)/len(last_10):.3f}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 11 — Save + Push to HF Hub"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.save_pretrained('outputs/lora_adapter')\n",
-    "tokenizer.save_pretrained('outputs/lora_adapter')\n",
-    "print('✅ LoRA adapter saved')\n",
-    "\n",
-    "print('Merging adapter into base model (bf16)...')\n",
-    "merged_model, merged_tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name='outputs/lora_adapter',\n",
-    "    load_in_4bit=False,\n",
-    "    max_seq_length=2048,\n",
-    ")\n",
-    "merged_model.save_pretrained_merged(\n",
-    "    'outputs/merged',\n",
-    "    merged_tokenizer,\n",
-    "    save_method='merged_16bit',\n",
-    ")\n",
-    "print('✅ Merged model saved')\n",
-    "\n",
-    "if HF_REPO and HF_TOKEN:\n",
-    "    print(f'Pushing to {HF_REPO}...')\n",
-    "    merged_model.push_to_hub_merged(\n",
-    "        HF_REPO,\n",
-    "        merged_tokenizer,\n",
-    "        save_method='merged_16bit',\n",
-    "        token=HF_TOKEN,\n",
-    "    )\n",
-    "    print(f'✅ Model live at https://huggingface.co/{HF_REPO}')\n",
-    "else:\n",
-    "    print('⚠️  Set HF_REPO and HF_TOKEN to push to Hub')\n",
-    "\n",
-    "print('Done.')"
-   ]
-  }
- ]
-}
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

grpo_train.py

@@ -13,6 +13,16 @@ from trl import GRPOTrainer, GRPOConfig
 
 PatchFastRL("GRPO", FastLanguageModel)
 
+# #region agent log
+import pathlib as _pl
+_DLOG = _pl.Path("debug-851b5f.log")
+def _dlog(hyp, loc, msg, data=None):
+    import time as _t
+    entry = json.dumps({"sessionId":"851b5f","hypothesisId":hyp,"location":loc,"message":msg,"data":data or {},"timestamp":int(_t.time()*1000)})
+    with open(_DLOG, "a") as f: f.write(entry + "\n")
+    print(f"[DBG:{hyp}] {msg} {data or ''}", flush=True)
+# #endregion
+
 # =========================
 # CONFIG
 # =========================
@@ -37,7 +47,10 @@ ALLOWED_ACTIONS = [
 # =========================
 
 def ensure_env_ready():
-    for _ in range(20):
+    # #region agent log
+    _dlog("B", "grpo_train.py:ensure_env_ready", "Checking env", {"ENV_URL": ENV_URL})
+    # #endregion
+    for i in range(20):
         try:
             r = requests.post(
                 f"{ENV_URL}/reset",
@@ -45,11 +58,20 @@ def ensure_env_ready():
                 timeout=5
             )
             if r.status_code == 200:
+                # #region agent log
+                _dlog("B", "grpo_train.py:ensure_env_ready", "Env ready", {"attempt": i+1, "status": r.status_code})
+                # #endregion
                 print("✅ Environment ready")
                 return
-        except:
+        except Exception as e:
+            # #region agent log
+            if i == 0: _dlog("B", "grpo_train.py:ensure_env_ready", "Env connection failed", {"attempt": i+1, "error": str(e)[:200]})
+            # #endregion
             pass
         time.sleep(1)
+    # #region agent log
+    _dlog("B", "grpo_train.py:ensure_env_ready", "ENV UNREACHABLE after 20 attempts", {})
+    # #endregion
     raise RuntimeError("❌ ENV not reachable")
 
 # =========================
@@ -240,21 +262,39 @@ def build_dataset():
 # REWARD FUNCTION (FIXED)
 # =========================
 
+_reward_call_count = [0]
+
 def reward_environment(prompts, completions, task_id=None, setup_actions=None, **kwargs):
-    """Shaped reward for GRPO.
-
-    Pure env reward is too sparse (mostly -0.05) to give clear gradients.
-    We add explicit shaping:
-      - invalid JSON / invalid action_type -> -1.0  (strong negative signal)
-      - valid action env REJECTS (wrong phase / API failure) -> -0.5
-      - valid action env ACCEPTS (advances state) -> +0.5 + env_reward
-      - terminal correct decision -> env_reward already contains +1.0 bonus
-    """
+    """Shaped reward for GRPO."""
+    _reward_call_count[0] += 1
+    _call = _reward_call_count[0]
+    # #region agent log
+    _dlog("C", "grpo_train.py:reward_env", f"reward call #{_call}", {
+        "n_prompts": len(prompts) if prompts else 0,
+        "n_completions": len(completions) if completions else 0,
+        "completions_type": type(completions).__name__,
+        "first_completion_type": type(completions[0]).__name__ if completions else "N/A",
+        "first_completion_preview": str(completions[0])[:150] if completions else "N/A",
+        "task_id_is_none": task_id is None,
+        "setup_actions_is_none": setup_actions is None,
+        "kwargs_keys": list(kwargs.keys()),
+    })
+    # #endregion
+
     client = EnvClient(ENV_URL)
     rewards = []
 
-    for completion, t_id, setup in zip(completions, task_id, setup_actions):
+    if task_id is None or setup_actions is None:
+        # #region agent log
+        _dlog("D", "grpo_train.py:reward_env", "task_id or setup_actions is None — returning -1 for all", {"call": _call})
+        # #endregion
+        return [-1.0] * len(completions)
+
+    for idx, (completion, t_id, setup) in enumerate(zip(completions, task_id, setup_actions)):
         parsed = extract_json(completion)
+        # #region agent log
+        if _call <= 3: _dlog("D", "grpo_train.py:reward_loop", f"call#{_call} item#{idx}", {"parsed_ok": parsed is not None, "action": parsed.get("action_type") if parsed else None, "raw_preview": str(completion)[:120], "task_id": t_id})
+        # #endregion
         if not parsed:
             rewards.append(-1.0)
             continue
@@ -300,23 +340,39 @@ def reward_environment(prompts, completions, task_id=None, setup_actions=None, *
 # MODEL
 # =========================
 
-USE_4BIT = not torch.cuda.is_available() or torch.cuda.get_device_properties(0).total_memory < 40 * 1024**3
+if torch.cuda.is_available():
+    _props = torch.cuda.get_device_properties(0)
+    _vram = _props.total_memory
+    _name = _props.name
+    _cc = (_props.major, _props.minor)  # compute capability
+    print(f"GPU: {_name}  VRAM: {_vram / 1024**3:.1f} GB  Compute: {_cc[0]}.{_cc[1]}")
+else:
+    _vram = 0
+    _name = "CPU"
+    _cc = (0, 0)
+
+USE_4BIT = _vram < 40 * 1024**3   # T4 (15 GB), L4 (24 GB) → 4-bit; A100 (80 GB) → full
+USE_BF16 = _cc >= (8, 0) and not USE_4BIT  # bf16 only when full-precision; 4-bit LoRA uses fp16 internally
+
+# #region agent log
+_dlog("A", "grpo_train.py:gpu_detect", "GPU config resolved", {"name":_name,"vram_gb":round(_vram/1024**3,1),"cc":list(_cc),"USE_4BIT":USE_4BIT,"USE_BF16":USE_BF16})
+# #endregion
 
 model, tokenizer = FastLanguageModel.from_pretrained(
     model_name="unsloth/Llama-3.1-8B-Instruct",
     load_in_4bit=USE_4BIT,
     max_seq_length=2048,
-    dtype=None,  # auto-detect bf16 on A100
+    dtype=torch.float16 if USE_4BIT else None,
 )
 
 model = FastLanguageModel.get_peft_model(
     model,
-    r=32,
+    r=16 if USE_4BIT else 32,
     target_modules=[
         "q_proj", "k_proj", "v_proj", "o_proj",
         "gate_proj", "up_proj", "down_proj",
     ],
-    lora_alpha=64,
+    lora_alpha=32 if USE_4BIT else 64,
     lora_dropout=0,
     bias="none",
     use_gradient_checkpointing="unsloth",
@@ -329,21 +385,26 @@ model = FastLanguageModel.get_peft_model(
 
 dataset = build_dataset()
 
+# #region agent log
+_dlog("A", "grpo_train.py:trainer_init", "Creating GRPOTrainer", {"USE_4BIT":USE_4BIT,"USE_BF16":USE_BF16,"epochs":1 if USE_4BIT else 3,"batch":1 if USE_4BIT else 2,"gens":2 if USE_4BIT else 4,"dataset_len":len(dataset)})
+# #endregion
+
 trainer = GRPOTrainer(
     model=model,
     reward_funcs=[reward_environment],
     args=GRPOConfig(
         output_dir="outputs",
         learning_rate=2e-5,
-        num_train_epochs=3,
-        per_device_train_batch_size=2,
-        gradient_accumulation_steps=4,
-        num_generations=4,
+        num_train_epochs=1 if USE_4BIT else 3,
+        per_device_train_batch_size=1 if USE_4BIT else 2,
+        gradient_accumulation_steps=2 if USE_4BIT else 4,
+        num_generations=2 if USE_4BIT else 4,
         max_prompt_length=768,
         max_completion_length=128,
-        logging_steps=5,
-        warmup_ratio=0.1,
-        bf16=True,
+        logging_steps=3 if USE_4BIT else 5,
+        warmup_steps=5 if USE_4BIT else 10,
+        bf16=USE_BF16,
+        fp16=not USE_BF16,
         report_to="none",
     ),
     train_dataset=dataset,
@@ -357,6 +418,9 @@ trainer = GRPOTrainer(
 if __name__ == "__main__":
     ensure_env_ready()
 
+    # #region agent log
+    _dlog("E", "grpo_train.py:train_start", "About to call trainer.train()", {"gpu_mem_allocated_gb": round(torch.cuda.memory_allocated()/1024**3, 2) if torch.cuda.is_available() else 0})
+    # #endregion
     print("Starting GRPO training...")
     trainer.train()