Update training/train_agent.py

training/train_agent.py

@@ -1,101 +1,100 @@
 """
 training/train_agent.py — SQL Database Engineer Agent
-Unsloth + GRPO training script.
-Run on venue GPU (April 25-26) with compute credits.
-
-FIXES applied:
-  1. Robust JSON extraction via regex (kills PARSE FALLBACK)
-  2. task_id from kwargs directly — not from kwargs["batch"] (kills only-easy_s001)
-  3. Reward calls /grader (stateless) instead of /reset+/step (kills race condition + flat 0.500)
-  4. Format bonus so valid JSON gets non-zero reward even before agent learns DBA actions
+Unsloth + GRPO — Fixed Version
+- 100 steps (not 700)
+- Local DatabaseSimulator for variance (not just /grader)
+- Real delta rewards, milestones shown in logs
+- Generates loss_curve.png automatically
 """
 
-import os
-import re
-import json
-import warnings
-import requests
-from datasets import Dataset
+import os, re, json, sys, warnings
+from pathlib import Path
 
-# ── Try importing Unsloth (GPU only) ─────────────────────────
+# ── Suppress known warnings ───────────────────────────────────
+warnings.filterwarnings("ignore", message=r".*max_new_tokens.*max_length.*")
+warnings.filterwarnings("ignore", message=r".*AttentionMaskConverter.*", category=FutureWarning)
+
+# ── GPU check + Unsloth ───────────────────────────────────────
+UNSLOTH_AVAILABLE = False
 try:
+    import torch
+    if not torch.cuda.is_available():
+        print("❌ No GPU. Unsloth requires CUDA.")
+        sys.exit(1)
     from unsloth import FastLanguageModel
     from trl import GRPOTrainer, GRPOConfig
+    from datasets import Dataset
     UNSLOTH_AVAILABLE = True
-except ImportError:
-    UNSLOTH_AVAILABLE = False
-    print("⚠️  Unsloth not available. Run: pip install unsloth trl")
-
-# Keep demo/training logs readable by silencing known upstream deprecation noise.
-warnings.filterwarnings(
-    "ignore",
-    message=r".*Both `max_new_tokens`.*and `max_length`.*",
-)
-warnings.filterwarnings(
-    "ignore",
-    message=r".*AttentionMaskConverter.*deprecated.*",
-    category=FutureWarning,
-)
-warnings.filterwarnings(
-    "ignore",
-    message=r".*use_return_dict.*deprecated.*",
-)
-
-# ─────────────────────────────────────────────
-#  CONFIG
-# ─────────────────────────────────────────────
-
+    print(f"✅ GPU: {torch.cuda.get_device_name(0)}")
+    print(f"✅ VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.1f}GB")
+except ImportError as e:
+    print(f"❌ {e}")
+    sys.exit(1)
+
+# ── Add project root to path ──────────────────────────────────
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+from env.db_simulator import DatabaseSimulator
+
+# ── Config ────────────────────────────────────────────────────
 ENV_URL    = os.getenv("ENV_URL",    "https://junaid0600-sql-db-engineer-agent.hf.space")
 HF_TOKEN   = os.getenv("HF_TOKEN",  "")
-MODEL_NAME = os.getenv("MODEL_NAME", "unsloth/Qwen2.5-7B-Instruct")
+MODEL_NAME = os.getenv("MODEL_NAME", "unsloth/Qwen2.5-1.5B-Instruct")
 OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./sdea-trained")
+MAX_STEPS  = int(os.getenv("MAX_STEPS", "100"))
 
-# Valid Round 2 action types — model must use one of these
-VALID_ACTION_TYPES = {
+print(f"\n[CONFIG] Model:     {MODEL_NAME}")
+print(f"[CONFIG] Max steps: {MAX_STEPS}")
+print(f"[CONFIG] Output:    {OUTPUT_DIR}\n")
+
+VALID_ACTIONS = {
     "inspect_query", "analyze_indexes", "create_index",
     "rewrite_query", "add_column", "drop_index",
     "partition_table", "analyze_statistics",
     "request_hint", "submit_report",
 }
 
-SYSTEM_PROMPT = """You are a senior database engineer.
-Given a database scenario with slow queries, choose the BEST single action to improve performance.
+SYSTEM_PROMPT = """You are a senior database engineer fixing slow database queries.
 
-Investigation pattern (follow this order):
-1. Use inspect_query to understand WHY a query is slow (scan type, rows examined)
-2. Use analyze_indexes to see what indexes exist and what is missing
-3. Use create_index to add the missing index on WHERE/JOIN columns
-4. Use rewrite_query if the SQL itself is inefficient
-5. Use partition_table for tables with 1M+ rows and range queries
-6. Use submit_report when performance target is reached
+Investigation pattern:
+1. inspect_query  → understand WHY query is slow
+2. analyze_indexes → see what indexes are missing
+3. create_index   → add composite index on WHERE/JOIN columns
+4. submit_report  → when performance target is reached
 
-RESPOND WITH VALID JSON ONLY. No explanation. No markdown. No preamble.
+RESPOND WITH VALID JSON ONLY. No markdown. No explanation.
 Examples:
 {"action_type": "inspect_query", "payload": {"query_id": "q1"}}
-{"action_type": "analyze_indexes", "payload": {"table": "users"}}
 {"action_type": "create_index", "payload": {"table": "users", "columns": ["email"]}}
-{"action_type": "create_index", "payload": {"table": "orders", "columns": ["user_id", "status"]}}
-{"action_type": "submit_report", "payload": {"summary": "Added composite index on orders(user_id, status). Performance improved from 5.0 to 85.0."}}"""
+{"action_type": "create_index", "payload": {"table": "orders", "columns": ["user_id", "status"]}}"""
 
 
-# ─────────────────────────────────────────────
-#  JSON EXTRACTION  (FIX 1 — kills PARSE FALLBACK)
-# ─────────────────────────────────────────────
+# ── Load all 15 scenarios ─────────────────────────────────────
+def load_scenarios() -> list:
+    scenarios = []
+    for fname in ["easy_scenarios.json", "medium_scenarios.json", "hard_scenarios.json"]:
+        path = os.path.join(ROOT, "dataset", fname)
+        try:
+            with open(path) as f:
+                data = json.load(f)
+                scenarios.extend(data)
+                print(f"  ✅ {len(data)} from {fname}")
+        except FileNotFoundError:
+            print(f"  ⚠️  {fname} not found")
+    print(f"  Total: {len(scenarios)} scenarios\n")
+    return scenarios
 
-def _extract_json(text: str) -> dict | None:
-    """
-    Robustly extract a JSON object from model output.
-    Handles: pure JSON, markdown blocks, JSON buried in text, partial JSON.
-    Returns parsed dict or None if nothing parseable found.
-    """
+ALL_SCENARIOS = load_scenarios()
+
+
+# ── Parse LLM output → action dict ───────────────────────────
+def parse_action(text: str) -> dict | None:
     if not text:
         return None
-
-    # Strip common markdown wrappers
     text = text.strip()
     text = re.sub(r"```(?:json)?", "", text).replace("```", "").strip()
 
-    # Try 1: entire text is valid JSON
+    # Try full parse
     try:
         obj = json.loads(text)
         if isinstance(obj, dict) and "action_type" in obj:
@@ -103,334 +102,296 @@ def _extract_json(text: str) -> dict | None:
     except json.JSONDecodeError:
         pass
 
-    # Try 2: balanced-brace scan from every opening '{'
-    # This handles nested payload objects/arrays more reliably than regex.
-    starts = [idx for idx, ch in enumerate(text) if ch == "{"]
-    for start in starts:
-        depth = 0
-        in_str = False
-        escape = False
+    # Try extract from partial text
+    for start in [i for i,c in enumerate(text) if c == "{"]:
+        depth, in_str, escape = 0, False, False
         for i in range(start, len(text)):
             ch = text[i]
             if in_str:
-                if escape:
-                    escape = False
-                elif ch == "\\":
-                    escape = True
-                elif ch == '"':
-                    in_str = False
+                if escape: escape = False
+                elif ch == "\\": escape = True
+                elif ch == '"': in_str = False
                 continue
-            if ch == '"':
-                in_str = True
-                continue
-            if ch == "{":
-                depth += 1
+            if ch == '"': in_str = True; continue
+            if ch == "{": depth += 1
             elif ch == "}":
                 depth -= 1
                 if depth == 0:
-                    candidate = text[start:i + 1]
                     try:
-                        obj = json.loads(candidate)
+                        obj = json.loads(text[start:i+1])
                         if isinstance(obj, dict) and "action_type" in obj:
                             return obj
-                    except json.JSONDecodeError:
-                        pass
+                    except: pass
                     break
-
     return None
 
 
-def _is_valid_action(action: dict) -> bool:
-    """Check action has correct structure before sending to /grader."""
-    if not isinstance(action, dict):
-        return False
-    if "action_type" not in action:
-        return False
-    if action["action_type"] not in VALID_ACTION_TYPES:
-        return False
-    if "payload" not in action or not isinstance(action.get("payload"), dict):
-        return False
-    return True
-
-
-def _completion_to_text(completion) -> str:
+# ── Local reward using DatabaseSimulator ──────────────────────
+def compute_reward(action: dict, scenario: dict) -> tuple:
     """
-    Normalize TRL completion formats into raw text.
-    Handles dict/list/string and content blocks emitted by some chat templates.
+    Compute reward LOCALLY — no HTTP, no shared state, deterministic.
+    Returns (score, improvement_pts, milestone_bonus, description)
     """
-    if isinstance(completion, str):
-        return completion
-
-    if isinstance(completion, dict):
-        content = completion.get("content", "")
-        if isinstance(content, str):
-            return content
-        if isinstance(content, list):
-            chunks = []
-            for block in content:
-                if isinstance(block, dict):
-                    chunks.append(str(block.get("text", "")))
-                else:
-                    chunks.append(str(block))
-            return "".join(chunks)
-        return str(content)
-
-    if isinstance(completion, list):
-        # Common shape: [{"role":"assistant","content":"..."}]
-        if not completion:
-            return ""
-        return "".join(_completion_to_text(item) for item in completion)
-
-    return str(completion)
-
-
-# ─────────────────────────────────────────────
-#  REWARD FUNCTION  (FIX 2 + FIX 3)
-# ─────────────────────────────────────────────
-
+    sim      = DatabaseSimulator(scenario)
+    baseline = sim.get_performance_score()
+    action_type = action.get("action_type", "")
+    payload     = action.get("payload", {})
+
+    # Apply action
+    delta = 0.0
+    if action_type == "create_index":
+        result = sim.apply_action("create_index", payload)
+        delta  = result.get("delta", 0.0)
+    elif action_type == "rewrite_query":
+        result = sim.apply_action("rewrite_query", payload)
+        delta  = result.get("delta", 0.0)
+    elif action_type == "partition_table":
+        result = sim.apply_action("partition_table", payload)
+        delta  = result.get("delta", 0.0)
+    elif action_type == "analyze_statistics":
+        result = sim.apply_action("analyze_statistics", payload)
+        delta  = result.get("delta", 0.0)
+    elif action_type in ("inspect_query", "analyze_indexes"):
+        delta = 0.0  # investigation — no DB change
+    elif action_type == "submit_report":
+        delta = max(0, sim.get_performance_score() - baseline)
+
+    final       = sim.get_performance_score()
+    improvement = max(0.0, final - baseline)
+    max_possible = max(1.0, 100.0 - baseline)
+    ratio = improvement / max_possible
+
+    # Step reward
+    step_r = {
+        "inspect_query":     0.10,
+        "analyze_indexes":   0.10,
+        "create_index":      0.15,
+        "rewrite_query":     0.20,
+        "analyze_statistics":0.08,
+        "partition_table":   0.15,
+        "submit_report":     0.05,
+    }.get(action_type, 0.001)
+
+    # Delta reward — key signal
+    delta_r = min(0.65, ratio * 0.65)
+
+    # Milestone bonus
+    milestone = 0.0
+    milestone_str = ""
+    if ratio >= 0.75:
+        milestone = 0.40
+        milestone_str = "🎯 75% milestone!"
+    elif ratio >= 0.50:
+        milestone = 0.25
+        milestone_str = "🎯 50% milestone!"
+    elif ratio >= 0.25:
+        milestone = 0.15
+        milestone_str = "🎯 25% milestone!"
+
+    # Wrong index penalty
+    wrong_pen = -0.15 if (action_type == "create_index" and delta <= 0.0) else 0.0
+
+    total = max(0.001, min(0.999, step_r + delta_r + milestone + wrong_pen))
+    desc  = f"+{improvement:.1f}pts delta={delta:.1f} {milestone_str}"
+
+    return total, improvement, milestone, desc
+
+
+# ── GRPO reward function ──────────────────────────────────────
 def reward_fn(prompts, completions, **kwargs):
     """
-    GRPO reward function — calls /grader (STATELESS).
-
-    FIX 2: task_ids from kwargs["task_id"] directly (TRL passes dataset
-            columns as direct kwargs, NOT inside a "batch" key).
-
-    FIX 3: calls /grader instead of /reset + /step.
-            /grader is stateless — no race condition, no global env mutation,
-            no flat reward from concurrent resets overwriting each other.
+    LOCAL reward — DatabaseSimulator directly.
+    Rewards vary 0.001 to 0.999 giving GRPO real gradient signal.
     """
-    rewards  = []
-
-    # ── FIX 2: correct task_id extraction ────────────────────────
-    # TRL GRPO passes dataset columns directly as kwargs.
-    # With num_generations=4, each task_id is repeated 4x in the list.
-    raw_task_ids = kwargs.get("task_id", [])
-    if isinstance(raw_task_ids, str):
-        raw_task_ids = [raw_task_ids]
+    rewards = []
 
     for i, (prompt, completion) in enumerate(zip(prompts, completions)):
-        task_id = (
-            raw_task_ids[i]
-            if i < len(raw_task_ids)
-            else "easy_s001"
-        )
-
-        # ── Extract text from completion ──────────────────────────
-        text = _completion_to_text(completion)
+        try:
+            # Get text
+            if isinstance(completion, list):
+                text = completion[0].get("content","") if completion else ""
+            else:
+                text = str(completion)
+
+            # Pick scenario
+            scenario = ALL_SCENARIOS[i % len(ALL_SCENARIOS)]
+            sid      = scenario["id"]
+
+            # Parse
+            action = parse_action(text)
+
+            if action is None:
+                print(f"  [REWARD] {sid} | INVALID JSON | score=0.001")
+                rewards.append(0.001)
+                continue
 
-        # ── FIX 1: robust JSON parse ──────────────────────────────
-        action = _extract_json(text)
+            if action.get("action_type") not in VALID_ACTIONS:
+                print(f"  [REWARD] {sid} | UNKNOWN ACTION | score=0.05")
+                rewards.append(0.05)
+                continue
 
-        if action is None:
-            # Complete parse failure — 0.001 (not 0.0, avoids GRPO div-by-zero)
-            print(f"  [REWARD] scenario={task_id} | INVALID JSON | score=0.001")
-            rewards.append(0.001)
-            continue
-
-        # Format bonus: valid JSON with correct structure = small positive signal
-        # This gives the model SOMETHING to learn from even before it learns
-        # the right actions, avoiding the all-zero gradient problem.
-        if not _is_valid_action(action):
-            # JSON parsed but action_type is wrong/missing
-            print(f"  [REWARD] scenario={task_id} | INVALID ACTION | score=0.05")
-            rewards.append(0.05)
-            continue
-
-        # ── FIX 3: stateless /grader call ────────────────────────
-        try:
-            resp = requests.post(
-                f"{ENV_URL}/grader",
-                json={"task_id": task_id, "action": action},
-                timeout=20,
-            )
-            resp.raise_for_status()
-            score = float(resp.json().get("score", 0.001))
-            score = max(0.001, min(0.999, score))
-            print(f"  [REWARD] scenario={task_id} | action={action['action_type']} | score={score:.3f}")
+            # Compute locally
+            score, improvement, milestone, desc = compute_reward(action, scenario)
             rewards.append(score)
 
-        except requests.exceptions.Timeout:
-            print(f"  [REWARD] scenario={task_id} | GRADER TIMEOUT | score=0.05")
-            rewards.append(0.05)   # grader timed out — give format credit
+            print(f"  [REWARD] {sid} | "
+                  f"action={action['action_type']} | "
+                  f"{desc} | score={score:.3f}")
+
         except Exception as e:
-            print(f"[reward_fn] grader call failed for {task_id}: {e}")
+            print(f"  [REWARD] Error: {e}")
             rewards.append(0.001)
 
     return rewards
 
 
-# ─────────────────────────────────────────────
-#  BUILD TRAINING DATASET
-# ─────────────────────────────────────────────
-
-def build_dataset():
-    """
-    Build training examples from all Round 2 scenario JSON files.
-    Each example: {"prompt": "...", "task_id": "easy_s001"}.
-    task_id is passed through to reward_fn via kwargs (TRL behaviour).
-    """
-    scenarios = []
-
-    for fname in [
-        "dataset/easy_scenarios.json",
-        "dataset/medium_scenarios.json",
-        "dataset/hard_scenarios.json",
-    ]:
-        try:
-            with open(fname) as f:
-                loaded = json.load(f)
-                scenarios.extend(loaded)
-                print(f"  Loaded {len(loaded)} scenarios from {fname}")
-        except FileNotFoundError:
-            print(f"  {fname} not found, skipping")
-
-    if not scenarios:
-        print("  Falling back to /tasks endpoint...")
-        try:
-            resp     = requests.get(f"{ENV_URL}/tasks", timeout=15)
-            tasks    = resp.json().get("tasks", [])
-            scenarios = [{"id": t["id"], "description": t.get("description", "")}
-                         for t in tasks if "_s" in t["id"]]
-        except Exception as e:
-            print(f"  /tasks fallback failed: {e}")
-            # Minimal fallback so training doesn't crash
-            scenarios = [{"id": "easy_s001",
-                          "description": "User lookup query taking 2s. Add index.",
-                          "tables": [{"name": "users", "rows": 10000, "indexes": ["PRIMARY"]}],
-                          "slow_queries": [{"id": "q1", "sql": "SELECT * FROM users WHERE email=?", "avg_ms": 2000}],
-                          "performance_score_baseline": 8.0,
-                          "target_score": 80.0}]
-
+# ── Build dataset ─────────────────────────────────────────────
+def build_dataset() -> Dataset:
     examples = []
-    for s in scenarios:
-        tables_txt = json.dumps(s.get("tables", []), separators=(",", ":"))
-        queries_txt = json.dumps(s.get("slow_queries", []), separators=(",", ":"))
-        baseline    = s.get("performance_score_baseline", s.get("performance_score", 0))
-        target      = s.get("target_score", 85)
-        max_steps   = s.get("max_steps", 50)
-
+    for s in ALL_SCENARIOS:
+        tables_str  = json.dumps(s.get("tables", []))
+        queries_str = json.dumps(s.get("slow_queries", []))
+        hints_str   = json.dumps(s.get("missing_index_hints", []))
         prompt = (
-        f"{SYSTEM_PROMPT}\n\n"
-        f"Scenario: {s.get('id')} — {s.get('description','')}\n"
-        f"Slow queries: {', '.join(q['id'] for q in s.get('slow_queries',[]))}\n"
-        f"Tables: {', '.join(t['name'] for t in s.get('tables',[]))}\n"
-        f"Missing indexes: {json.dumps(s.get('missing_index_hints',[]))}\n"
-        f"Performance: {baseline}/100 → target {target}/100\n\n"
-        f"Respond with JSON only:"
-    )
-
-
-        examples.append({
-            "prompt":  prompt,
-            "task_id": s.get("id", "easy_s001"),
-        })
+            f"{SYSTEM_PROMPT}\n\n"
+            f"=== DATABASE STATE ===\n"
+            f"Scenario: {s['id']} — {s.get('description','')}\n"
+            f"Tables: {tables_str}\n"
+            f"Slow Queries: {queries_str}\n"
+            f"Missing Index Hints: {hints_str}\n"
+            f"Performance: {s.get('performance_score_baseline',0)}/100 "
+            f"→ Target: {s.get('target_score',85)}/100\n\n"
+            f"What is your next action? JSON only:"
+        )
+        examples.append({"prompt": prompt, "scenario_id": s["id"]})
 
-    print(f"Built {len(examples)} training examples from {len(scenarios)} scenarios")
+    print(f"  ✅ Dataset: {len(examples)} examples")
     return Dataset.from_list(examples)
 
 
-# ─────────────────────────────────────────────
-#  REWARD WRAPPER  (FIX 2 continued)
-# ─────────────────────────────────────────────
-
-def reward_wrapper(prompts, completions, **kwargs):
-    """
-    Thin wrapper — passes kwargs straight through.
-    TRL GRPO sends dataset columns (including task_id) as direct kwargs.
-    DO NOT use kwargs.get("batch") — that key does not exist in TRL GRPO.
-    """
-    return reward_fn(prompts, completions, **kwargs)
+# ── Generate loss + reward plots ──────────────────────────────
+def generate_plots(trainer):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
 
-
-# ─────────────────────────────────────────────
-#  MAIN TRAINING
-# ─────────────────────────────────────────────
-
-def train():
-    if not UNSLOTH_AVAILABLE:
-        print("Cannot train — Unsloth not installed")
-        print("Run: pip install 'unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git' trl transformers datasets accelerate")
+    logs = [l for l in trainer.state.log_history if "loss" in l]
+    if not logs:
+        print("⚠️ No logs to plot")
         return
 
-    print(f"🚀 Loading model: {MODEL_NAME}")
-    print(f"🌐 Environment:   {ENV_URL}")
-
-    # Sanity check — make sure environment is reachable
-    try:
-        r = requests.get(f"{ENV_URL}/health", timeout=10)
-        print(f"✅ Environment health: {r.json()}")
-    except Exception as e:
-        print(f"⚠️  Cannot reach environment at {ENV_URL}: {e}")
-        print("   Training will likely fail — check ENV_URL")
+    steps   = [l.get("step", i) for i,l in enumerate(logs)]
+    losses  = [l.get("loss", 0) for l in logs]
+    rewards = [l.get("reward", 0) for l in logs]
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
+    fig.suptitle("GRPO Training — SQL Database Engineer Agent",
+                 fontsize=13, fontweight="bold")
+
+    ax1.plot(steps, losses, "b-o", lw=2, ms=4, label="Loss")
+    ax1.set_xlabel("Training Step")
+    ax1.set_ylabel("Loss")
+    ax1.set_title("Training Loss (↓ = model learning)")
+    ax1.grid(True, alpha=0.3)
+    if losses:
+        ax1.annotate(f"Start: {losses[0]:.4f}", xy=(steps[0], losses[0]),
+                    xytext=(steps[0]+1, losses[0]*1.1), fontsize=8, color="red")
+        ax1.annotate(f"End: {losses[-1]:.4f}", xy=(steps[-1], losses[-1]),
+                    xytext=(steps[-1]-8, losses[-1]*1.15), fontsize=8, color="green")
+
+    ax2.plot(steps, rewards, "g-o", lw=2, ms=4, label="Avg Reward")
+    ax2.set_xlabel("Training Step")
+    ax2.set_ylabel("Reward")
+    ax2.set_title("Reward During Training (↑ = improving)")
+    ax2.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig("loss_curve.png", dpi=150, bbox_inches="tight")
+    print("✅ loss_curve.png saved")
+    if losses:
+        print(f"   Loss:   {losses[0]:.4f} → {losses[-1]:.4f}")
+    if rewards:
+        valid = [r for r in rewards if r > 0]
+        if valid:
+            print(f"   Reward: {valid[0]:.4f} → {valid[-1]:.4f}")
+
+
+# ── Main ──────────────────────────────────────────────────────
+def train():
+    if not ALL_SCENARIOS:
+        print("❌ No scenarios loaded. Check dataset/ folder.")
+        sys.exit(1)
 
-    # Load model with Unsloth 4-bit quantization
+    print(f"⏳ Loading {MODEL_NAME}...")
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name     = MODEL_NAME,
-        max_seq_length = 4096,
+        max_seq_length = 2048,
         load_in_4bit   = True,
         token          = HF_TOKEN or None,
     )
-
-    # Add LoRA adapters
     model = FastLanguageModel.get_peft_model(
         model,
-        r              = 16,
-        lora_alpha     = 16,
-        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
-                          "gate_proj", "up_proj", "down_proj"],
-        lora_dropout   = 0,
-        bias           = "none",
-        use_gradient_checkpointing = "unsloth",
+        r=16, lora_alpha=16,
+        target_modules=["q_proj","k_proj","v_proj","o_proj",
+                        "gate_proj","up_proj","down_proj"],
+        lora_dropout=0, bias="none",
+        use_gradient_checkpointing="unsloth",
+        random_state=42,
     )
+    print("✅ Model + LoRA ready\n")
 
-    # Build dataset
     dataset = build_dataset()
 
-    # GRPO config
-    # Avoid generation warnings from simultaneously using max_length and max_new_tokens.
-    if hasattr(model, "generation_config"):
-        model.generation_config.max_length = None
-    if hasattr(model, "config"):
-        # Some stacks still read generation defaults from model.config.
-        model.config.max_length = None
-
     config = GRPOConfig(
         output_dir                  = OUTPUT_DIR,
-        num_train_epochs            = 100,      # was 3  → now 100 → gives ~120 steps
-        per_device_train_batch_size = 1,        # was 2  → frees memory for 1.5B
-        gradient_accumulation_steps = 2,        # was 8  → was hiding steps, now shows them
-        learning_rate               = 2e-5,     # was 5e-5 → lower = more stable for small model
-        max_completion_length       = 120,
+        max_steps                   = MAX_STEPS,
+        per_device_train_batch_size = 1,
+        gradient_accumulation_steps = 2,
+        learning_rate               = 2e-5,
+        max_completion_length       = 150,
         num_generations             = 4,
-        logging_steps               = 1,        # was 10 → see every step
-        save_steps                  = 20,
-        warmup_steps                = 10,       # was warmup_ratio (deprecated)
+        temperature                 = 1.0,
+        logging_steps               = 1,
+        save_steps                  = 25,
+        save_total_limit            = 3,
+        warmup_steps                = 10,
         report_to                   = "none",
+        remove_unused_columns       = False,
     )
 
     trainer = GRPOTrainer(
         model         = model,
         tokenizer     = tokenizer,
-        reward_funcs  = reward_wrapper,
+        reward_funcs  = reward_fn,
         args          = config,
         train_dataset = dataset,
     )
 
-    print("🏋️  Starting GRPO training...")
-    print("   Expected reward progression:")
-    print("   Steps  10: ~0.05-0.15 (model still outputting free text)")
-    print("   Steps  50: ~0.20-0.35 (learning JSON format)")
-    print("   Steps 100: ~0.35-0.50 (learning correct action types)")
-    print("   Steps 200: ~0.55-0.70 (learning DBA investigation pattern)")
-    print("   Steps 300: ~0.70-0.82 (strategic multi-action planning)")
+    print(f"🏋️  GRPO training — {MAX_STEPS} steps")
+    print("Expected rewards:")
+    print("  inspect_query (always):        ~0.10")
+    print("  create_index (wrong columns):  ~0.001")
+    print("  create_index (right columns):  ~0.75-0.99")
+    print("  GRPO learns: right create_index >> everything else\n")
 
     trainer.train()
+    print("\n✅ Training complete!")
 
-    # Save
+    Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)
     model.save_pretrained(f"{OUTPUT_DIR}/final")
     tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
-    print(f"✅ Training complete. Model saved to {OUTPUT_DIR}/final")
+    print(f"✅ Saved to {OUTPUT_DIR}/final")
+
+    generate_plots(trainer)
+
+    print("\n" + "="*50)
+    print("NEXT:")
+    print("  python training/evaluate_agent.py")
+    print("  git add loss_curve.png reward_curve.png")
+    print("  git commit -m 'Real GRPO training evidence'")
+    print("  git push origin main")
+    print("="*50)
 
 
 if __name__ == "__main__":
-    train()
+    train()