Update training/train_agent.py

training/train_agent.py

@@ -126,7 +126,7 @@ def parse_action(text: str) -> dict | None:
     return None
 
 
-# ── Local reward using DatabaseSimulator ──────────────────────
+
 def compute_reward(action: dict, scenario: dict) -> tuple:
     """
     Compute reward LOCALLY — no HTTP, no shared state, deterministic.
@@ -142,6 +142,28 @@ def compute_reward(action: dict, scenario: dict) -> tuple:
     if action_type == "create_index":
         result = sim.apply_action("create_index", payload)
         delta  = result.get("delta", 0.0)
+
+        # ── Hint-match fallback ────────────────────────────────────────────
+        # When the simulator returns delta=0 (wrong/partial columns), check
+        # the scenario's missing_index_hints and simulate a proportional delta.
+        # This gives GRPO a gradient signal early in training before the model
+        # has learned the exact right column names.
+        if delta <= 0.0:
+            hints         = scenario.get("missing_index_hints", [])
+            p_table       = payload.get("table", "")
+            p_cols        = set(str(c).lower() for c in payload.get("columns", []))
+            best_match    = 0.0
+            for hint in hints:
+                h_table = hint.get("table", "")
+                h_cols  = set(str(c).lower() for c in hint.get("columns", []))
+                if not h_cols:
+                    continue
+                table_ok    = 1.0 if h_table == p_table else 0.3
+                col_overlap = len(p_cols & h_cols) / max(len(h_cols), 1)
+                best_match  = max(best_match, table_ok * 0.4 + col_overlap * 0.6)
+            if best_match > 0:
+                max_gap = max(1.0, 100.0 - baseline)
+                delta   = best_match * max_gap * 0.65   # up to 65% of gap on perfect match → reaches 25% milestone
     elif action_type == "rewrite_query":
         result = sim.apply_action("rewrite_query", payload)
         delta  = result.get("delta", 0.0)
@@ -156,10 +178,12 @@ def compute_reward(action: dict, scenario: dict) -> tuple:
     elif action_type == "submit_report":
         delta = max(0, sim.get_performance_score() - baseline)
 
-    final       = sim.get_performance_score()
-    improvement = max(0.0, final - baseline)
+    final        = sim.get_performance_score()
+    sim_improve  = max(0.0, final - baseline)
+    # Use hint-simulated delta if simulator returned 0 (wrong columns early in training)
+    improvement  = max(sim_improve, delta)
     max_possible = max(1.0, 100.0 - baseline)
-    ratio = improvement / max_possible
+    ratio        = improvement / max_possible
 
     # Step reward
     step_r = {
@@ -175,24 +199,28 @@ def compute_reward(action: dict, scenario: dict) -> tuple:
     # Delta reward — key signal
     delta_r = min(0.65, ratio * 0.65)
 
-    # Milestone bonus
-    milestone = 0.0
-    milestone_str = ""
+    # Milestone bonus — cumulative, all thresholds crossed shown + rewarded
+    milestone       = 0.0
+    milestone_parts = []
+    if ratio >= 0.25:
+        milestone += 0.15
+        milestone_parts.append("25%")
+    if ratio >= 0.50:
+        milestone += 0.25
+        milestone_parts.append("50%")
     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
+        milestone += 0.40
+        milestone_parts.append("75%")
+    milestone_str = ("🎯 " + "+".join(milestone_parts) + " milestone!") if milestone_parts else ""
+
+    # Wrong index penalty reduced: -0.05 instead of -0.15
+    # Old value exactly cancelled step_r (0.15 - 0.15 = 0 → clamped to 0.001)
+    # Now wrong create_index still scores ~0.10 instead of 0.001
+    wrong_pen = -0.05 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}"
+    src   = "sim" if sim_improve > 0 else ("hint" if delta > 0 else "none")
+    desc  = f"+{improvement:.1f}pts delta={delta:.1f}[{src}] {milestone_str}"
 
     return total, improvement, milestone, desc
 
@@ -270,7 +298,14 @@ def build_dataset() -> Dataset:
 
 
 # ── Generate loss + reward plots ──────────────────────────────
+# ──────────────────────────────────────────────────────────────
+# REPLACE ONLY THIS FUNCTION in your existing train_agent.py
+# Find: def generate_plots(trainer):
+# Replace the entire function with this:
+# ──────────────────────────────────────────────────────────────
+
 def generate_plots(trainer):
+    import json, numpy as np
     import matplotlib
     matplotlib.use("Agg")
     import matplotlib.pyplot as plt
@@ -280,40 +315,77 @@ def generate_plots(trainer):
         print("⚠️ No logs to plot")
         return
 
-    steps   = [l.get("step", i) for i,l in enumerate(logs)]
-    losses  = [l.get("loss", 0) for l in logs]
+    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",
+    # Save logs for generate_plots.py
+    import os, json
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    with open(f"{OUTPUT_DIR}/training_logs.json", "w") as f:
+        json.dump(trainer.state.log_history, f)
+    print(f"✅ Logs saved to {OUTPUT_DIR}/training_logs.json")
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 5))
+    fig.suptitle("GRPO Training — SQL Database Engineer Agent\n"
+                 "Qwen2.5-1.5B fine-tuned with Unsloth + TRL",
                  fontsize=13, fontweight="bold")
 
-    ax1.plot(steps, losses, "b-o", lw=2, ms=4, label="Loss")
+    # ── LEFT: Loss with smoothing ──────────────────────────────
+    ax1.plot(steps, losses, "b-", lw=1.0, alpha=0.35, label="Raw loss")
+    if len(losses) >= 10:
+        smooth = np.convolve(losses, np.ones(10)/10, mode="valid")
+        ax1.plot(steps[9:], smooth, "b-", lw=2.5, label="10-step avg")
     ax1.set_xlabel("Training Step")
     ax1.set_ylabel("Loss")
-    ax1.set_title("Training Loss (↓ = model learning)")
+    ax1.set_title("Training Loss  ↓ = model learning")
+    ax1.ticklabel_format(style="sci", axis="y", scilimits=(0,0))
     ax1.grid(True, alpha=0.3)
+    ax1.legend(fontsize=9)
+    # FIX 1: scientific notation annotations
     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")
+        ax1.annotate(f"Start: {losses[0]:.2e}",
+                     xy=(steps[0], losses[0]),
+                     xytext=(steps[0]+3, max(losses)*0.85),
+                     fontsize=8, color="red",
+                     arrowprops=dict(arrowstyle="->", color="red", lw=1))
+        ax1.annotate(f"End: {losses[-1]:.2e}",
+                     xy=(steps[-1], losses[-1]),
+                     xytext=(steps[-1]-12, max(losses)*0.65),
+                     fontsize=8, color="green",
+                     arrowprops=dict(arrowstyle="->", color="green", lw=1))
+
+    # ── RIGHT: Reward with smoothing ───────────────────────────
+    ax2.plot(steps, rewards, "g-", lw=1.0, alpha=0.35, label="Raw reward")
+    if len(rewards) >= 10:
+        smooth_r = np.convolve(rewards, np.ones(10)/10, mode="valid")
+        ax2.plot(steps[9:], smooth_r, "g-", lw=2.5, label="10-step avg")
     ax2.set_xlabel("Training Step")
-    ax2.set_ylabel("Reward")
-    ax2.set_title("Reward During Training (↑ = improving)")
+    ax2.set_ylabel("Avg Reward")
+    ax2.set_title("Reward During Training  ↑ = improving")
     ax2.grid(True, alpha=0.3)
-
-    plt.tight_layout()
+    ax2.legend(fontsize=9)
+
+    # Bottom summary
+    if losses and rewards:
+        start_r = rewards[0]
+        end_r   = rewards[-1]
+        pct     = ((end_r-start_r)/max(abs(start_r),1e-9))*100
+        fig.text(0.5, 0.01,
+                 f"Loss: {losses[0]:.2e} → {losses[-1]:.2e}  |  "
+                 f"Reward: {start_r:.3f} → {end_r:.3f} ({'+'if pct>=0 else ''}{pct:.0f}%)",
+                 ha="center", fontsize=10,
+                 bbox=dict(boxstyle="round", facecolor="lightyellow", alpha=0.8))
+
+    plt.tight_layout(rect=[0, 0.07, 1, 1])
     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}")
+        print(f"   Loss:   {losses[0]:.2e} → {losses[-1]:.2e}")
     if rewards:
         valid = [r for r in rewards if r > 0]
         if valid:
-            print(f"   Reward: {valid[0]:.4f} → {valid[-1]:.4f}")
+            print(f"   Reward: {valid[0]:.3f} → {valid[-1]:.3f}")
 
 
 # ── Main ──────────────────────────────────────────────────────
@@ -369,10 +441,12 @@ def train():
 
     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")
+    print("  inspect_query / analyze_indexes:       ~0.10")
+    print("  create_index (no table/col match):     ~0.10  (was 0.001)")
+    print("  create_index (partial hint match):     ~0.20-0.45")
+    print("  create_index (perfect hint match):     ~0.55-0.80")
+    print("  create_index (simulator confirms):     ~0.75-0.99")
+    print("  Milestones: 25%=+0.15  50%=+0.25  75%=+0.40 (cumulative)\n")
 
     trainer.train()
     print("\n✅ Training complete!")
@@ -394,4 +468,4 @@ def train():
 
 
 if __name__ == "__main__":
-    train()
+    train()