Upload routing agent adapters and source code

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Deterministic_Action_Space_GRPO.pdf filter=lfs diff=lfs merge=lfs -text

Deterministic_Action_Space_GRPO.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3671c1a7060dd74216760d76d7fe15649addd9c725df5e821c5317126fa8d403
+size 1003880

adapters/router_grpo_epoch_1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82defe5a768a836545176fc4c19759ff4f5de39fc31e18bbac2d35053f9952b2
+size 741423

adapters/router_grpo_epoch_10.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4fde8546c2beb8f9c0af61feafd7835417549433d31b9b628de67a8642e08e7
+size 741593

adapters/router_grpo_epoch_11.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ee3dcc8dcc885a4cdb538878e398b300681bd7c7e5f26bfd697bf47302b2750
+size 741593

adapters/router_grpo_epoch_12.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f15932966200969a4950c4b8157f67ce06e81db36062fd2dde512b50f8562e48
+size 741593

adapters/router_grpo_epoch_13.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ea6d2cbbd6371e603a292606a82d7dfab3c5a2a23cacb09c083018bcab0c8f5
+size 741593

adapters/router_grpo_epoch_14.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a4680d4923c2b55ddc8905a9048fba7fd0a66de585e6a09fa06be645532f6ca
+size 741593

adapters/router_grpo_epoch_15.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d8ceb15f07681860b398181a20c82f02a4dacd8cbcb53061d74ee643949c30b0
+size 741593

adapters/router_grpo_epoch_16.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d16c1b921a4c39b36266f1ffab0ff5be0bf1b14cc9e9d1577db73c3d830d00c
+size 741593

adapters/router_grpo_epoch_17.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2958e857269bdffb584d99aebd71f7e3ec9a1539bcbb4cebabb268ecb9d2804
+size 741593

adapters/router_grpo_epoch_18.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebb8538a30f10af1f1236f955ae31e590737fcc6a0ee0f6362e5a2bcd668878e
+size 741593

adapters/router_grpo_epoch_19.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6a1f132eb63e286ad598e0875e543b5c277b6c5063ae3ab3509bec2508f1bfe
+size 741593

adapters/router_grpo_epoch_2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ac6f190372b0e60167d15a7a1b57c5a1501ff8110359d1f1d6c566b5aa2c15ac
+size 741423

adapters/router_grpo_epoch_20.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:740562659427d570ebffed75ca38994722042b487b8cd1d65a8fa606449b8d6f
+size 741593

adapters/router_grpo_epoch_3.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5d90746091add53d56728c0e7ed42e66febf858dc445b0859569ec1be21d5903
+size 741423

adapters/router_grpo_epoch_4.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dcbfb3f72c8853bb4e99669e58a37c855d85555c674afc5b79f8f3b1c00332d
+size 741423

adapters/router_grpo_epoch_5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0f4c6f4e1c929d77932244105065f30090eacff5b8f36e98ba7a0aa6ac7d0421
+size 741423

adapters/router_grpo_epoch_6.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a3afa0288f4374e62aa8fd11899afd41987958e61760dfc60cd6f018ed14b58
+size 741551

adapters/router_grpo_epoch_7.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fb00b41a4151b0f02797bf659f439d690388fc6922b96a14001efedb46a6df0c
+size 741551

adapters/router_grpo_epoch_8.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4846278c2b838dd126063cde2af02acd2ca826ea78e517455f37e27a0ab7f907
+size 741551

adapters/router_grpo_epoch_9.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:04cf97b6a22b1954c0020e2e482a47a6d931c935f311ce906e7fbfb477e8954e
+size 741551

benchmark_evaluation.py

@@ -0,0 +1,735 @@
+#!/usr/bin/env python3
+"""
+benchmark_evaluation.py — Comparative evaluation suite for RL router vs baseline.
+
+Evaluates two routing approaches on held-out evaluation sets:
+  1. Baseline: backend's own LLM router (Qwen3-14B) — no search_sources passed
+  2. GRPO: local RL router — explicit search_sources passed to backend
+
+For each query the suite:
+  - Routes via the assigned strategy
+  - Executes the backend RAG pipeline
+  - Compares predicted action to ground-truth action_label
+  - Records accuracy, latency, and raw outputs
+
+Outputs:
+  - Per-strategy detailed CSV results
+  - Aggregated comparison tables (CSV + LaTeX-ready)
+  - Statistical summaries
+"""
+
+import os
+import sys
+import time
+import json
+import csv
+import uuid
+import logging
+import argparse
+import statistics
+from datetime import datetime
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional
+
+# ---------------------------------------------------------------------------
+# Path / env setup
+# ---------------------------------------------------------------------------
+ROOT = Path(__file__).resolve().parent
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+REPO_ROOT = ROOT.parent
+sys.path.append(str(REPO_ROOT))
+sys.path.append(str(REPO_ROOT / "backend"))
+
+import torch
+import pandas as pd
+import requests
+from sentence_transformers import SentenceTransformer
+from scipy import stats as scipy_stats
+
+from model import RouterNetwork
+
+# ---------------------------------------------------------------------------
+# Logging
+# ---------------------------------------------------------------------------
+logger = logging.getLogger("benchmark_eval")
+if not logger.handlers:
+    logger.setLevel(logging.INFO)
+    fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
+    ch = logging.StreamHandler()
+    ch.setFormatter(fmt)
+    logger.addHandler(ch)
+    logs_dir = REPO_ROOT / "logs"
+    try:
+        os.makedirs(logs_dir, exist_ok=True)
+        fh = logging.FileHandler(logs_dir / "benchmark_evaluation.log")
+        fh.setFormatter(fmt)
+        logger.addHandler(fh)
+    except Exception:
+        pass
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+NUM_ACTIONS = 3
+ACTION_TO_TOOL = {0: "guidelines", 1: "web-search", 2: "course-registration"}
+ACTION_TO_NAME = {0: "Guidelines", 1: "Web-Search", 2: "CBNU Database"}
+
+BACKEND_CHAT_URL = os.getenv("BACKEND_CHAT_URL", "http://localhost:2000/api/chat")
+
+BASELINE_LABEL = "Baseline-LLM"
+
+TOP_GRPO_CHECKPOINTS = [
+    {
+        "id": "GRPO-E14",
+        "label": "GRPO-E14",
+        "path": ROOT / "adapters" / "router_grpo_epoch_14.pt",
+        "epoch": 14,
+        "train_avg_reward": 0.4262,
+    },
+    {
+        "id": "GRPO-E12",
+        "label": "GRPO-E12",
+        "path": ROOT / "adapters" / "router_grpo_epoch_12.pt",
+        "epoch": 12,
+        "train_avg_reward": 0.4189,
+    },
+]
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def detect_language(text: str) -> str:
+    if any("\uac00" <= char <= "\ud7a3" for char in text):
+        return "ko"
+    elif any("\u4e00" <= char <= "\u9fff" for char in text):
+        return "zh"
+    return "en"
+
+
+def truncate(s: str, n: int = 500) -> str:
+    if s is None:
+        return ""
+    s = str(s)
+    return s if len(s) <= n else s[: n - 1] + "\u2026"
+
+
+# ---------------------------------------------------------------------------
+# Backend RAG pipeline
+# ---------------------------------------------------------------------------
+def simulate_rag_pipeline(query: str, search_sources: Optional[List[str]] = None, max_retries: int = 3) -> Tuple[str, str]:
+    """Call the backend chat API.
+
+    If search_sources is None, the backend's own LLM router decides which
+    tool(s) to use (baseline behavior).
+    If search_sources is provided, the backend is forced to use those tools
+    (GRPO behavior).
+
+    Includes automatic retry with exponential backoff on 429 responses.
+    """
+    url = BACKEND_CHAT_URL
+    headers = {"accept": "application/json", "Content-Type": "application/json"}
+    language = detect_language(query)
+
+    payload = {
+        "message": query,
+        "language": language,
+        "session_id": f"benchmark-{uuid.uuid4().hex}",
+    }
+    if search_sources is not None:
+        payload["search_sources"] = search_sources
+
+    last_exception = None
+    for attempt in range(max_retries):
+        try:
+            response = requests.post(url, headers=headers, json=payload, timeout=60)
+            if response.status_code == 429:
+                wait_time = (2 ** attempt) * 2  # 2s, 4s, 8s
+                logger.warning(f"Rate limited (429). Retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
+                time.sleep(wait_time)
+                continue
+            response.raise_for_status()
+            data = response.json()
+            answer = data.get("message", "")
+            context_parts = []
+            for src in data.get("sources", []):
+                if isinstance(src, dict) and src.get("success"):
+                    for res in src.get("data", {}).get("results", []):
+                        if "content" in res:
+                            context_parts.append(res["content"])
+                        elif "snippet" in res:
+                            context_parts.append(res["snippet"])
+            context = "\n".join(context_parts)
+            return context, answer
+        except Exception as e:
+            last_exception = e
+            if attempt < max_retries - 1:
+                wait_time = (2 ** attempt) * 2
+                logger.warning(f"Request failed: {e}. Retrying in {wait_time}s...")
+                time.sleep(wait_time)
+            else:
+                break
+
+    logger.error(f"simulate_rag_pipeline failed after {max_retries} attempts: {last_exception}")
+    return "", ""
+
+
+# ---------------------------------------------------------------------------
+# GRPO router inference
+# ---------------------------------------------------------------------------
+def load_grpo_checkpoint(path: Path, device: str = "cpu"):
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    embed_dim = embedder.get_sentence_embedding_dimension()
+    model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)
+
+    sd = torch.load(str(path), map_location="cpu")
+    if isinstance(sd, dict) and "model_state_dict" in sd:
+        model.load_state_dict(sd["model_state_dict"], strict=False)
+    else:
+        model.load_state_dict(sd, strict=False)
+    model.eval()
+    return model, embedder
+
+
+def grpo_route_query(model: RouterNetwork, embedder: SentenceTransformer, query: str, device: str = "cpu") -> int:
+    """Route query using GRPO model, return action_idx."""
+    with torch.no_grad():
+        emb = embedder.encode(query, convert_to_tensor=True).to(device)
+        action, _log_prob, _value = model.get_action(emb)
+    return int(action.item())
+
+
+def action_to_search_sources(action_idx: int) -> List[str]:
+    if action_idx == 0:
+        return ["guidelines"]
+    elif action_idx == 1:
+        return ["web-search"]
+    else:
+        return ["course-registration", "department-info"]
+
+
+# ---------------------------------------------------------------------------
+# Evaluation runner
+# ---------------------------------------------------------------------------
+def evaluate_strategy(
+    name: str,
+    queries: List[str],
+    ground_truths: List[int],
+    route_fn,
+    call_tool: bool = True,
+    force_search_sources: bool = True,
+    delay: float = 1.0,
+) -> pd.DataFrame:
+    """Run a routing strategy over all queries and return a results DataFrame.
+
+    Args:
+        force_search_sources: If True, pass explicit search_sources derived from
+            the routed action_idx. If False, pass search_sources=None so the
+            backend's own router decides (baseline behavior).
+        delay: Delay in seconds between consecutive API calls to avoid rate limiting.
+    """
+    results = []
+    for idx, (q, gt) in enumerate(zip(queries, ground_truths)):
+        row = {"question": q, "ground_truth": gt}
+        try:
+            # Routing decision
+            t0 = time.perf_counter()
+            action_idx = route_fn(q)
+            t1 = time.perf_counter()
+            row["predicted_label"] = action_idx
+            row["routing_time_sec"] = t1 - t0
+            row["routing_error"] = ""
+
+            # Tool execution
+            context, answer = "", ""
+            if call_tool:
+                search_sources = action_to_search_sources(action_idx) if force_search_sources else None
+                t2 = time.perf_counter()
+                context, answer = simulate_rag_pipeline(q, search_sources=search_sources)
+                t3 = time.perf_counter()
+                row["tool_time_sec"] = t3 - t2
+                row["end_to_end_time_sec"] = t3 - t0
+                row["context"] = truncate(context, 1000)
+                row["answer"] = truncate(answer, 1000)
+            else:
+                row["tool_time_sec"] = None
+                row["end_to_end_time_sec"] = row["routing_time_sec"]
+                row["context"] = ""
+                row["answer"] = ""
+        except Exception as e:
+            logger.error(f"Error on query: {truncate(q, 80)} — {e}")
+            row["routing_error"] = str(e)
+            row["predicted_label"] = None
+            row["routing_time_sec"] = None
+            row["tool_time_sec"] = None
+            row["end_to_end_time_sec"] = None
+            row["context"] = ""
+            row["answer"] = ""
+
+        row["correct"] = 1 if row.get("predicted_label") == gt else 0
+        results.append(row)
+
+        # Delay between queries to avoid rate limiting
+        if delay > 0 and idx < len(queries) - 1:
+            time.sleep(delay)
+
+    return pd.DataFrame(results)
+
+
+# ---------------------------------------------------------------------------
+# Metrics computation
+# ---------------------------------------------------------------------------
+def compute_metrics(df: pd.DataFrame) -> Dict:
+    """Compute accuracy, latency stats from a results DataFrame."""
+    valid = df.dropna(subset=["predicted_label", "ground_truth"])
+    accuracy = float((valid["predicted_label"] == valid["ground_truth"]).mean()) if len(valid) else 0.0
+
+    e2e_valid = df.dropna(subset=["end_to_end_time_sec"])
+    e2e_mean = float(e2e_valid["end_to_end_time_sec"].mean()) if len(e2e_valid) else 0.0
+    e2e_std = float(e2e_valid["end_to_end_time_sec"].std()) if len(e2e_valid) > 1 else 0.0
+
+    routing_valid = df.dropna(subset=["routing_time_sec"])
+    routing_mean = float(routing_valid["routing_time_sec"].mean()) if len(routing_valid) else 0.0
+    routing_std = float(routing_valid["routing_time_sec"].std()) if len(routing_valid) > 1 else 0.0
+
+    tool_valid = df.dropna(subset=["tool_time_sec"])
+    tool_mean = float(tool_valid["tool_time_sec"].mean()) if len(tool_valid) else 0.0
+    tool_std = float(tool_valid["tool_time_sec"].std()) if len(tool_valid) > 1 else 0.0
+
+    # Per-class precision / recall / F1
+    classes = sorted(valid["ground_truth"].unique().tolist())
+    per_class = {}
+    for cls in classes:
+        tp = int(((valid["predicted_label"] == cls) & (valid["ground_truth"] == cls)).sum())
+        fp = int(((valid["predicted_label"] == cls) & (valid["ground_truth"] != cls)).sum())
+        fn = int(((valid["predicted_label"] != cls) & (valid["ground_truth"] == cls)).sum())
+        prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0
+        rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+        f1 = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0.0
+        per_class[int(cls)] = {"precision": prec, "recall": rec, "f1": f1, "support": int((valid["ground_truth"] == cls).sum())}
+
+    # Confusion matrix
+    conf_mat = {}
+    for gt_cls in classes:
+        for pred_cls in classes:
+            conf_mat[(int(gt_cls), int(pred_cls))] = int(
+                ((valid["ground_truth"] == gt_cls) & (valid["predicted_label"] == pred_cls)).sum()
+            )
+
+    # Action distribution
+    action_dist = valid["predicted_label"].value_counts().to_dict()
+    action_dist = {int(k): int(v) for k, v in action_dist.items()}
+
+    return {
+        "num_queries": int(len(df)),
+        "num_valid": int(len(valid)),
+        "num_errors": int(len(df) - len(valid)),
+        "accuracy": accuracy,
+        "e2e_latency_mean": e2e_mean,
+        "e2e_latency_std": e2e_std,
+        "routing_latency_mean": routing_mean,
+        "routing_latency_std": routing_std,
+        "tool_latency_mean": tool_mean,
+        "tool_latency_std": tool_std,
+        "per_class": per_class,
+        "confusion_matrix": conf_mat,
+        "action_distribution": action_dist,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Statistical tests
+# ---------------------------------------------------------------------------
+def paired_t_test_latency(baseline_latencies: List[float], treatment_latencies: List[float]) -> Dict:
+    """Paired t-test between baseline and treatment latency vectors."""
+    if len(baseline_latencies) != len(treatment_latencies) or len(baseline_latencies) < 2:
+        return {"t_stat": None, "p_value": None, "significant": None}
+
+    t_stat, p_value = scipy_stats.ttest_rel(treatment_latencies, baseline_latencies)
+    return {
+        "t_stat": float(t_stat),
+        "p_value": float(p_value),
+        "significant": p_value < 0.05,
+    }
+
+
+def mean_confidence_interval(data: List[float], confidence: float = 0.95) -> Tuple[float, float, float]:
+    """Return (mean, lower, upper) for CI."""
+    if not data:
+        return 0.0, 0.0, 0.0
+    n = len(data)
+    mean = statistics.mean(data)
+    if n < 2:
+        return mean, mean, mean
+    se = statistics.stdev(data) / (n ** 0.5)
+    h = se * scipy_stats.t.ppf((1 + confidence) / 2.0, n - 1)
+    return mean, mean - h, mean + h
+
+
+# ---------------------------------------------------------------------------
+# Table / report generation
+# ---------------------------------------------------------------------------
+def build_comparison_table(all_metrics: Dict[str, Dict], split_name: str) -> str:
+    """Build a Markdown comparison table for a dataset split."""
+    rows = []
+    rows.append(f"### {split_name} Evaluation Results\n")
+    rows.append("| Metric | Baseline-LLM | GRPO-E14 | GRPO-E12 |")
+    rows.append("|--------|-------------|---------|---------|")
+
+    rows.append(
+        f"| Routing Accuracy | {all_metrics['Baseline-LLM']['accuracy']:.2%} | "
+        f"{all_metrics['GRPO-E14']['accuracy']:.2%} | {all_metrics['GRPO-E12']['accuracy']:.2%} |"
+    )
+    rows.append(
+        f"| E2E Latency (mean ± std, s) | {all_metrics['Baseline-LLM']['e2e_latency_mean']:.4f} ± {all_metrics['Baseline-LLM']['e2e_latency_std']:.4f} | "
+        f"{all_metrics['GRPO-E14']['e2e_latency_mean']:.4f} ± {all_metrics['GRPO-E14']['e2e_latency_std']:.4f} | "
+        f"{all_metrics['GRPO-E12']['e2e_latency_mean']:.4f} ± {all_metrics['GRPO-E12']['e2e_latency_std']:.4f} |"
+    )
+    rows.append(
+        f"| Routing Latency (mean ± std, s) | {all_metrics['Baseline-LLM']['routing_latency_mean']:.4f} ± {all_metrics['Baseline-LLM']['routing_latency_std']:.4f} | "
+        f"{all_metrics['GRPO-E14']['routing_latency_mean']:.4f} ± {all_metrics['GRPO-E14']['routing_latency_std']:.4f} | "
+        f"{all_metrics['GRPO-E12']['routing_latency_mean']:.4f} ± {all_metrics['GRPO-E12']['routing_latency_std']:.4f} |"
+    )
+    rows.append(
+        f"| Tool Latency (mean ± std, s) | {all_metrics['Baseline-LLM']['tool_latency_mean']:.4f} ± {all_metrics['Baseline-LLM']['tool_latency_std']:.4f} | "
+        f"{all_metrics['GRPO-E14']['tool_latency_mean']:.4f} ± {all_metrics['GRPO-E14']['tool_latency_std']:.4f} | "
+        f"{all_metrics['GRPO-E12']['tool_latency_mean']:.4f} ± {all_metrics['GRPO-E12']['tool_latency_std']:.4f} |"
+    )
+    rows.append("")
+    return "\n".join(rows)
+
+
+def build_latex_table(all_metrics: Dict[str, Dict], split_name: str) -> str:
+    """Build a LaTeX table suitable for direct inclusion in a journal paper."""
+    def bold_if_best(vals: List[float], idx: int, fmt_fn, higher_is_better: bool = True) -> str:
+        best_idx = 0
+        best_val = vals[0]
+        for i, v in enumerate(vals[1:], 1):
+            if (higher_is_better and v > best_val) or (not higher_is_better and v < best_val):
+                best_val = v
+                best_idx = i
+        if idx == best_idx:
+            return "\\textbf{" + fmt_fn(vals[idx]) + "}"
+        return fmt_fn(vals[idx])
+
+    b = all_metrics["Baseline-LLM"]
+    e14 = all_metrics["GRPO-E14"]
+    e12 = all_metrics["GRPO-E12"]
+
+    accs = [b["accuracy"], e14["accuracy"], e12["accuracy"]]
+    e2es = [b["e2e_latency_mean"], e14["e2e_latency_mean"], e12["e2e_latency_mean"]]
+    routings = [b["routing_latency_mean"], e14["routing_latency_mean"], e12["routing_latency_mean"]]
+    tools = [b["tool_latency_mean"], e14["tool_latency_mean"], e12["tool_latency_mean"]]
+
+    rows = []
+    rows.append("% Auto-generated by benchmark_evaluation.py")
+    rows.append("\\begin{table}[ht]")
+    rows.append("\\centering")
+    rows.append(f"\\caption{{Performance comparison on {split_name} evaluation set. "
+                f"Results reported as mean $\\pm$ standard deviation. "
+                f"Bold indicates best result.}}")
+    rows.append("\\label{tab:comparison_" + split_name.lower() + "}")
+    rows.append("\\begin{tabular}{lccc}")
+    rows.append("\\toprule")
+    rows.append("Metric & Baseline-LLM & GRPO-E14 & GRPO-E12 \\\\")
+    rows.append("\\midrule")
+
+    def acc_fmt(v): return f"{v:.2%}"
+    def lat_fmt(v): return f"{v:.3f}"
+
+    rows.append(
+        f"Routing Accuracy & {bold_if_best(accs, 0, acc_fmt, higher_is_better=True)} & "
+        f"{bold_if_best(accs, 1, acc_fmt, higher_is_better=True)} & "
+        f"{bold_if_best(accs, 2, acc_fmt, higher_is_better=True)} \\\\"
+    )
+    rows.append(
+        f"E2E Latency (s) & {bold_if_best(e2es, 0, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(b['e2e_latency_std'])} & "
+        f"{bold_if_best(e2es, 1, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e14['e2e_latency_std'])} & "
+        f"{bold_if_best(e2es, 2, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e12['e2e_latency_std'])} \\\\"
+    )
+    rows.append(
+        f"Routing Latency (s) & {bold_if_best(routings, 0, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(b['routing_latency_std'])} & "
+        f"{bold_if_best(routings, 1, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e14['routing_latency_std'])} & "
+        f"{bold_if_best(routings, 2, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e12['routing_latency_std'])} \\\\"
+    )
+    rows.append(
+        f"Tool Latency (s) & {bold_if_best(tools, 0, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(b['tool_latency_std'])} & "
+        f"{bold_if_best(tools, 1, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e14['tool_latency_std'])} & "
+        f"{bold_if_best(tools, 2, lat_fmt, higher_is_better=False)}$\\pm${lat_fmt(e12['tool_latency_std'])} \\\\"
+    )
+
+    rows.append("\\bottomrule")
+    rows.append("\\end{tabular}")
+    rows.append("\\end{table}")
+    rows.append("")
+    return "\n".join(rows)
+
+
+def build_confusion_table(conf_mat: Dict, split_name: str) -> str:
+    """Build a confusion matrix table in Markdown."""
+    classes = sorted({k[0] for k in conf_mat.keys()})
+    rows = []
+    rows.append(f"#### Confusion Matrix — {split_name}\n")
+    header = "| Pred \\ GT | " + " | ".join(f"Class {c}" for c in classes) + " |"
+    rows.append(header)
+    sep = "|" + "|".join(["---"] * (len(classes) + 1)) + "|"
+    rows.append(sep)
+    for gt in classes:
+        vals = [str(conf_mat.get((gt, pred), 0)) for pred in classes]
+        rows.append(f"| Class {gt} | " + " | ".join(vals) + " |")
+    rows.append("")
+    return "\n".join(rows)
+
+
+def build_per_class_table(per_class: Dict, split_name: str) -> str:
+    """Build per-class precision/recall/F1 table."""
+    rows = []
+    rows.append(f"#### Per-Class Metrics — {split_name}\n")
+    rows.append("| Class | Support | Precision | Recall | F1 |")
+    rows.append("|-------|---------|-----------|--------|----|")
+    for cls_id in sorted(per_class.keys()):
+        m = per_class[cls_id]
+        rows.append(
+            f"| Class {cls_id} | {m['support']} | {m['precision']:.4f} | {m['recall']:.4f} | {m['f1']:.4f} |"
+        )
+    rows.append("")
+    return "\n".join(rows)
+
+
+def build_action_distribution_table(action_dist: Dict, split_name: str) -> str:
+    """Build action distribution table."""
+    rows = []
+    rows.append(f"#### Action Distribution — {split_name}\n")
+    rows.append("| Action | Count | Percentage |")
+    rows.append("|--------|-------|------------|")
+    total = sum(action_dist.values()) if action_dist else 0
+    for action_id in sorted(action_dist.keys()):
+        count = action_dist[action_id]
+        pct = count / total if total > 0 else 0.0
+        rows.append(f"| {ACTION_TO_NAME.get(action_id, f'Action {action_id}')} | {count} | {pct:.2%} |")
+    rows.append("")
+    return "\n".join(rows)
+
+
+def _stringify_keys(obj):
+    """Recursively convert dict keys to strings for JSON serialization."""
+    if isinstance(obj, dict):
+        return {str(k): _stringify_keys(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [_stringify_keys(item) for item in obj]
+    return obj
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    parser = argparse.ArgumentParser(description="Benchmark evaluation for RL router")
+    parser.add_argument("--splits", nargs="+", default=["eng", "ko"], choices=["eng", "ko", "both"],
+                        help="Which evaluation splits to run")
+    parser.add_argument("--output-dir", type=str, default=str(ROOT / "benchmark_results"),
+                        help="Directory for outputs")
+    parser.add_argument("--call-tool", action="store_true", default=True,
+                        help="Whether to call RAG tools (required for latency/accuracy measurement)")
+    parser.add_argument("--device", type=str, default="cpu",
+                        help="Device for GRPO model inference")
+    parser.add_argument("--max-queries", type=int, default=None,
+                        help="Limit queries per split (for smoke testing)")
+    parser.add_argument("--delay", type=float, default=1.0,
+                        help="Delay in seconds between API calls to avoid rate limiting")
+    parser.add_argument("--backoff-on-429", action="store_true", default=True,
+                        help="Enable exponential backoff when receiving 429 responses")
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir)
+    os.makedirs(output_dir, exist_ok=True)
+
+    split_files = {
+        "eng": REPO_ROOT / "routing_agent" / "grpo_benchmark_en.csv",
+        "ko": REPO_ROOT / "routing_agent" / "grpo_benchmark_ko.csv",
+    }
+
+    splits_to_run = args.splits if "both" not in args.splits else ["eng", "ko"]
+    all_split_metrics = {}
+    all_detailed_dfs = {}
+
+    strategies = [BASELINE_LABEL] + [ckpt["label"] for ckpt in TOP_GRPO_CHECKPOINTS]
+
+    for split in splits_to_run:
+        csv_path = split_files[split]
+        if not csv_path.exists():
+            logger.error(f"Evaluation file not found: {csv_path}")
+            continue
+
+        df_eval = pd.read_csv(csv_path)
+        queries = df_eval["question"].tolist()
+        ground_truths = df_eval["action_label"].astype(int).tolist()
+
+        if args.max_queries:
+            queries = queries[: args.max_queries]
+            ground_truths = ground_truths[: args.max_queries]
+
+        logger.info(f"=== Evaluating {split} split: {len(queries)} queries ===")
+
+        # Pre-load GRPO models
+        grpo_models = {}
+        for ckpt in TOP_GRPO_CHECKPOINTS:
+            logger.info(f"Loading GRPO checkpoint: {ckpt['label']} from {ckpt['path']}")
+            model, embedder = load_grpo_checkpoint(ckpt["path"], device=args.device)
+            grpo_models[ckpt["label"]] = (model, embedder)
+
+        split_results = {}
+
+        # ---- Baseline: backend LLM router, no search_sources ----
+        logger.info(f"Running {BASELINE_LABEL} on {split}...")
+        def baseline_route(q):
+            # For baseline, we don't need local routing — backend decides.
+            # Return a dummy action; the actual routing happens inside backend.
+            return 0
+
+        df_base = evaluate_strategy(
+            name=BASELINE_LABEL,
+            queries=queries,
+            ground_truths=ground_truths,
+            route_fn=baseline_route,
+            call_tool=args.call_tool,
+            force_search_sources=False,
+            delay=args.delay,
+        )
+        split_results[BASELINE_LABEL] = df_base
+
+        # ---- GRPO checkpoints ----
+        for ckpt in TOP_GRPO_CHECKPOINTS:
+            label = ckpt["label"]
+            model, embedder = grpo_models[label]
+            logger.info(f"Running {label} on {split}...")
+
+            def make_route_fn(m, e):
+                def fn(q):
+                    return grpo_route_query(m, e, q, device=args.device)
+                return fn
+
+            df_grpo = evaluate_strategy(
+                name=label,
+                queries=queries,
+                ground_truths=ground_truths,
+                route_fn=make_route_fn(model, embedder),
+                call_tool=args.call_tool,
+                force_search_sources=True,
+                delay=args.delay,
+            )
+            split_results[label] = df_grpo
+
+        # Save detailed CSV per strategy
+        for strategy, df_res in split_results.items():
+            out_csv = output_dir / f"{split}_{strategy}_detailed.csv"
+            df_res.to_csv(out_csv, index=False)
+            logger.info(f"Saved detailed results: {out_csv}")
+
+        # Compute metrics
+        metrics = {}
+        for strategy, df_res in split_results.items():
+            metrics[strategy] = compute_metrics(df_res)
+
+        # Statistical tests (baseline vs each GRPO) on E2E latency
+        base_latencies = split_results[BASELINE_LABEL]["end_to_end_time_sec"].dropna().tolist()
+        stat_tests = {}
+        for ckpt in TOP_GRPO_CHECKPOINTS:
+            label = ckpt["label"]
+            treat_latencies = split_results[label]["end_to_end_time_sec"].dropna().tolist()
+            stat_tests[label] = {
+                "paired_t": paired_t_test_latency(base_latencies, treat_latencies),
+                "baseline_latency_ci": mean_confidence_interval(base_latencies),
+                "treatment_latency_ci": mean_confidence_interval(treat_latencies),
+            }
+
+        all_split_metrics[split] = {"metrics": metrics, "stat_tests": stat_tests}
+        all_detailed_dfs[split] = split_results
+
+        # Print tables to console and save to markdown
+        md_tables = []
+        md_tables.append(build_comparison_table(metrics, split))
+        for strategy in [BASELINE_LABEL] + [ckpt["label"] for ckpt in TOP_GRPO_CHECKPOINTS]:
+            md_tables.append(build_confusion_table(metrics[strategy]["confusion_matrix"], split))
+            md_tables.append(build_per_class_table(metrics[strategy]["per_class"], split))
+            md_tables.append(build_action_distribution_table(metrics[strategy]["action_distribution"], split))
+
+        md_report = "\n".join(md_tables)
+        report_path = output_dir / f"{split}_report.md"
+        with open(report_path, "w", encoding="utf-8") as f:
+            f.write(md_report)
+        logger.info(f"Saved report: {report_path}")
+        print(md_report)
+
+        # Save LaTeX tables
+        latex_path = output_dir / f"{split}_tables.tex"
+        with open(latex_path, "w", encoding="utf-8") as f:
+            f.write(build_latex_table(metrics, split))
+            f.write("\n")
+        logger.info(f"Saved LaTeX tables: {latex_path}")
+
+        # Save JSON metrics for programmatic access
+        json_path = output_dir / f"{split}_metrics.json"
+        with open(json_path, "w", encoding="utf-8") as f:
+            json.dump(
+                _stringify_keys({
+                    "metrics": metrics,
+                    "stat_tests": {
+                        k: {
+                            t_label: (
+                                {tk: tv for tk, tv in t_val.items()}
+                                if isinstance(t_val, dict)
+                                else list(t_val)
+                            )
+                            for t_label, t_val in stat_tests[k].items()
+                        }
+                        for k, t_val in stat_tests.items()
+                    },
+                }),
+                f,
+                indent=2,
+                default=str,
+            )
+        logger.info(f"Saved JSON metrics: {json_path}")
+
+    # ---- Combined (both splits) aggregate ----
+    if len(splits_to_run) == 2:
+        logger.info("=== Building combined aggregate report ===")
+        combined_results = {}
+        for strategy in strategies:
+            combined_results[strategy] = pd.concat(
+                [all_detailed_dfs["eng"][strategy], all_detailed_dfs["ko"][strategy]], ignore_index=True
+            )
+            out_csv = output_dir / f"combined_{strategy}_detailed.csv"
+            combined_results[strategy].to_csv(out_csv, index=False)
+
+        combined_metrics = {s: compute_metrics(df) for s, df in combined_results.items()}
+
+        base_latencies_combined = combined_results[BASELINE_LABEL]["end_to_end_time_sec"].dropna().tolist()
+        combined_stat = {}
+        for ckpt in TOP_GRPO_CHECKPOINTS:
+            label = ckpt["label"]
+            treat = combined_results[label]["end_to_end_time_sec"].dropna().tolist()
+            combined_stat[label] = {
+                "paired_t": paired_t_test_latency(base_latencies_combined, treat),
+            }
+
+        md_combined = []
+        md_combined.append(build_comparison_table(combined_metrics, "Combined"))
+        report_path = output_dir / "combined_report.md"
+        with open(report_path, "w", encoding="utf-8") as f:
+            f.write("\n".join(md_combined))
+        logger.info(f"Saved combined report: {report_path}")
+        print("\n".join(md_combined))
+
+        latex_combined = output_dir / "combined_tables.tex"
+        with open(latex_combined, "w", encoding="utf-8") as f:
+            f.write(build_latex_table(combined_metrics, "Combined"))
+        logger.info(f"Saved combined LaTeX: {latex_combined}")
+
+    logger.info("Benchmark evaluation complete.")
+
+
+if __name__ == "__main__":
+    main()

eval_with_timing.py

@@ -0,0 +1,52 @@
+import os
+import sys
+import time
+import traceback
+import pandas as pd
+
+# Ensure local package imports work
+ROOT = os.path.dirname(__file__)
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+import run_inference as ri
+
+
+def predict_and_time(input_csv: str, output_csv: str):
+    df = pd.read_csv(input_csv)
+    results = []
+
+    for idx, row in df.iterrows():
+        q = row.get('question')
+        ground = row.get('action_label')
+        try:
+            start = time.perf_counter()
+            # fast routing-only prediction using loaded embedder + model
+            with ri.torch.no_grad():
+                emb = ri.embedder.encode(q, convert_to_tensor=True).to(ri.DEVICE)
+                action, log_prob, _ = ri.model.get_action(emb)
+                pred_idx = int(action.item())
+            end = time.perf_counter()
+            elapsed = end - start
+        except Exception:
+            pred_idx = None
+            elapsed = None
+            traceback.print_exc()
+
+        results.append({
+            'question': q,
+            'ground_truth': ground,
+            'predicted_label': pred_idx,
+            'time_took': elapsed,
+        })
+
+    out_df = pd.DataFrame(results)
+    out_df.to_csv(output_csv, index=False)
+    print(f"Wrote results to {output_csv} ({len(out_df)} rows)")
+
+
+if __name__ == '__main__':
+    inp = os.path.join(ROOT, 'eval_queries.csv')
+    out = os.path.join(ROOT, 'eval_with_timing_results.csv')
+    print(f"Input: {inp}\nOutput: {out}\nDevice: {ri.DEVICE}")
+    predict_and_time(inp, out)

generate_paper_results.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""
+generate_paper_results.py — Generate publication-ready results from benchmark outputs.
+
+Reads the per-split metrics JSON files produced by benchmark_evaluation.py and
+writes:
+  - LaTeX tables (main results, per-class, latency comparisons)
+  - Statistical significance summary
+  - Plain-text summaries
+
+Intended to be run AFTER benchmark_evaluation.py has completed.
+"""
+
+import json
+import os
+import logging
+import argparse
+from pathlib import Path
+from typing import Dict, List
+
+logger = logging.getLogger("generate_paper_results")
+
+STRATEGIES = ["Baseline-LLM", "GRPO-E14", "GRPO-E12"]
+
+
+def load_json(path: Path) -> Dict:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def fmt_pct(v: float) -> str:
+    return f"{v:.2%}"
+
+
+def fmt_num(v: float, decimals: int = 4) -> str:
+    return f"{v:.{decimals}f}"
+
+
+def best(values: List[float], higher_is_better: bool = True) -> int:
+    idx = 0
+    best_val = values[0]
+    for i, v in enumerate(values[1:], 1):
+        if (higher_is_better and v > best_val) or (not higher_is_better and v < best_val):
+            best_val = v
+            idx = i
+    return idx
+
+
+# ---------------------------------------------------------------------------
+# Table generators
+# ---------------------------------------------------------------------------
+def main_results_table(all_metrics: Dict[str, Dict]) -> str:
+    """Generate the primary results table for the Results section."""
+    lines = []
+    lines.append("% Auto-generated by generate_paper_results.py")
+    lines.append("\\begin{table}[ht]")
+    lines.append("\\centering")
+    lines.append("\\caption{Comparative evaluation of routing strategies. "
+                 "Mean $\\pm$ standard deviation reported. "
+                 "Best result per row is bolded.}")
+    lines.append("\\label{tab:main_results}")
+    lines.append("\\begin{tabular}{lccc}")
+    lines.append("\\toprule")
+    lines.append("Metric & Baseline-LLM (Qwen3-14B) & GRPO-E14 & GRPO-E12 \\\\")
+    lines.append("\\midrule")
+
+    # Routing Accuracy
+    accs = [all_metrics[s]["accuracy"] for s in STRATEGIES]
+    bi = best(accs, higher_is_better=True)
+    acc_strs = [fmt_pct(v) for v in accs]
+    acc_strs[bi] = "\\textbf{" + acc_strs[bi] + "}"
+    lines.append(f"Routing Accuracy & {acc_strs[0]} & {acc_strs[1]} & {acc_strs[2]} \\\\")
+
+    # E2E Latency
+    e2e_means = [all_metrics[s]["e2e_latency_mean"] for s in STRATEGIES]
+    e2e_stds = [all_metrics[s]["e2e_latency_std"] for s in STRATEGIES]
+    bi = best(e2e_means, higher_is_better=False)
+    cells = []
+    for i, s in enumerate(STRATEGIES):
+        cell = f"{fmt_num(e2e_means[i], 3)}$\\pm${fmt_num(e2e_stds[i], 3)}"
+        if i == bi:
+            cell = "\\textbf{" + cell + "}"
+        cells.append(cell)
+    lines.append(f"E2E Latency (s) & {cells[0]} & {cells[1]} & {cells[2]} \\\\")
+
+    # Routing Latency
+    r_means = [all_metrics[s]["routing_latency_mean"] for s in STRATEGIES]
+    r_stds = [all_metrics[s]["routing_latency_std"] for s in STRATEGIES]
+    bi = best(r_means, higher_is_better=False)
+    cells = []
+    for i, s in enumerate(STRATEGIES):
+        cell = f"{fmt_num(r_means[i], 4)}$\\pm${fmt_num(r_stds[i], 4)}"
+        if i == bi:
+            cell = "\\textbf{" + cell + "}"
+        cells.append(cell)
+    lines.append(f"Routing Latency (s) & {cells[0]} & {cells[1]} & {cells[2]} \\\\")
+
+    # Tool Latency
+    t_means = [all_metrics[s]["tool_latency_mean"] for s in STRATEGIES]
+    t_stds = [all_metrics[s]["tool_latency_std"] for s in STRATEGIES]
+    bi = best(t_means, higher_is_better=False)
+    cells = []
+    for i, s in enumerate(STRATEGIES):
+        cell = f"{fmt_num(t_means[i], 3)}$\\pm${fmt_num(t_stds[i], 3)}"
+        if i == bi:
+            cell = "\\textbf{" + cell + "}"
+        cells.append(cell)
+    lines.append(f"Tool Latency (s) & {cells[0]} & {cells[1]} & {cells[2]} \\\\")
+
+    lines.append("\\bottomrule")
+    lines.append("\\end{tabular}")
+    lines.append("\\end{table}")
+    lines.append("")
+    return "\n".join(lines)
+
+
+def per_class_table(all_metrics: Dict[str, Dict]) -> str:
+    """Generate per-class precision/recall/F1 table."""
+    lines = []
+    lines.append("\\begin{table}[ht]")
+    lines.append("\\centering")
+    lines.append("\\caption{Per-class routing performance.}")
+    lines.append("\\label{tab:per_class}")
+    lines.append("\\begin{tabular}{lcccccc}")
+    lines.append("\\toprule")
+    lines.append("Strategy & Class & Support & Precision & Recall & F1 \\\\")
+    lines.append("\\midrule")
+    for s in STRATEGIES:
+        pc = all_metrics[s]["per_class"]
+        for cls_id in sorted(pc.keys()):
+            m = pc[cls_id]
+            lines.append(
+                f"{s} & {cls_id} & {m['support']} & "
+                f"{fmt_num(m['precision'])} & {fmt_num(m['recall'])} & {fmt_num(m['f1'])} \\\\"
+            )
+    lines.append("\\bottomrule")
+    lines.append("\\end{tabular}")
+    lines.append("\\end{table}")
+    lines.append("")
+    return "\n".join(lines)
+
+
+def statistical_summary_table(stat_tests: Dict) -> str:
+    """Generate statistical significance summary."""
+    lines = []
+    lines.append("\\begin{table}[ht]")
+    lines.append("\\centering")
+    lines.append("\\caption{Paired statistical tests comparing each GRPO variant against the baseline. "
+                 "$t$-statistic and $p$-value from paired $t$-test on per-query end-to-end latency.}")
+    lines.append("\\label{tab:statistical_tests}")
+    lines.append("\\begin{tabular}{lcc}")
+    lines.append("\\toprule")
+    lines.append("Comparison & $t$-statistic & $p$-value \\\\")
+    lines.append("\\midrule")
+    for label, tests in stat_tests.items():
+        t = tests.get("paired_t", {})
+        t_stat = t.get("t_stat")
+        p_val = t.get("p_value")
+        t_str = fmt_num(t_stat, 3) if t_stat is not None else "---"
+        p_str = fmt_num(p_val, 4) if p_val is not None else "---"
+        sig = " *" if p_val is not None and p_val < 0.05 else ""
+        lines.append(f"Baseline vs {label} & {t_str} & {p_str}{sig} \\\\")
+    lines.append("\\bottomrule")
+    lines.append("\\end{tabular}")
+    lines.append("\\end{table}")
+    lines.append("")
+    return "\n".join(lines)
+
+
+def training_curve_table() -> str:
+    """Generate a table of training checkpoint performance."""
+    import torch
+    adapters_dir = Path(__file__).resolve().parent / "routing_agent" / "adapters"
+    rows = []
+    rows.append("\\begin{table}[ht]")
+    rows.append("\\centering")
+    rows.append("\\caption{GRPO training checkpoint performance (training reward).}")
+    rows.append("\\label{tab:training_curve}")
+    rows.append("\\begin{tabular}{lccc}")
+    rows.append("\\toprule")
+    rows.append("Epoch & Avg Reward & Avg Loss & KL Divergence \\\\")
+    rows.append("\\midrule")
+
+    checkpoints = []
+    for f in sorted(adapters_dir.glob("router_grpo_epoch_*.pt")):
+        ckpt = torch.load(str(f), map_location="cpu")
+        epoch = ckpt.get("epoch", int(f.stem.split("_")[-1]))
+        checkpoints.append((epoch, ckpt.get("avg_reward", 0), ckpt.get("avg_loss", 0), ckpt.get("kl_div", 0)))
+
+    # Show top-5 and last 5
+    checkpoints.sort(key=lambda x: x[1], reverse=True)
+    top5 = checkpoints[:5]
+    bottom5 = checkpoints[-5:] if len(checkpoints) > 5 else []
+    shown = top5 + bottom5
+
+    for epoch, reward, loss, kl in shown:
+        marker = " *" if epoch in [12, 14] else ""
+        rows.append(f"{epoch} & {fmt_num(reward, 4)} & {fmt_num(loss, 4)} & {fmt_num(kl, 6)}{marker} \\\\")
+
+    rows.append("\\bottomrule")
+    rows.append("\\end{tabular}")
+    rows.append("\\end{table}")
+    rows.append("")
+    return "\n".join(rows)
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+def main():
+    parser = argparse.ArgumentParser(description="Generate paper-ready results from benchmark JSONs")
+    parser.add_argument("--input-dir", type=str, default="benchmark_results",
+                        help="Directory containing benchmark metrics JSONs")
+    parser.add_argument("--output-dir", type=str, default="paper_results",
+                        help="Output directory for LaTeX tables")
+    parser.add_argument("--splits", nargs="+", default=["eng", "ko", "combined"],
+                        help="Splits to include in aggregate")
+    args = parser.parse_args()
+
+    input_dir = Path(args.input_dir)
+    output_dir = Path(args.output_dir)
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Load all metrics
+    all_metrics = {}
+    all_stat_tests = {}
+    for split in args.splits:
+        json_path = input_dir / f"{split}_metrics.json"
+        if not json_path.exists():
+            logger.warning(f"Missing metrics file: {json_path}")
+            continue
+        data = load_json(json_path)
+        for s, m in data["metrics"].items():
+            if s not in all_metrics:
+                all_metrics[s] = {}
+            for k, v in m.items():
+                if k not in all_metrics[s]:
+                    all_metrics[s][k] = v
+                else:
+                    # average across splits
+                    if isinstance(v, (int, float)):
+                        all_metrics[s][k] = (all_metrics[s][k] + v) / 2
+        for s, st in data["stat_tests"].items():
+            if s not in all_stat_tests:
+                all_stat_tests[s] = st
+
+    if not all_metrics:
+        print("No metrics loaded. Run benchmark_evaluation.py first.")
+        return
+
+    # Generate LaTeX
+    latex_content = []
+    latex_content.append(main_results_table(all_metrics))
+    latex_content.append(per_class_table(all_metrics))
+    latex_content.append(statistical_summary_table(all_stat_tests))
+    latex_content.append(training_curve_table())
+
+    latex_path = output_dir / "results_tables.tex"
+    with open(latex_path, "w", encoding="utf-8") as f:
+        f.write("\n".join(latex_content))
+    print(f"Generated LaTeX tables: {latex_path}")
+
+    # Generate plain-text summary
+    txt_path = output_dir / "results_summary.txt"
+    with open(txt_path, "w", encoding="utf-8") as f:
+        f.write("=" * 70 + "\n")
+        f.write("BENCHMARK RESULTS SUMMARY\n")
+        f.write("=" * 70 + "\n\n")
+        for s in STRATEGIES:
+            m = all_metrics[s]
+            f.write(f"[{s}]\n")
+            f.write(f"  Accuracy:          {fmt_pct(m['accuracy'])}\n")
+            f.write(f"  E2E Latency (s):   {fmt_num(m['e2e_latency_mean'], 3)} ± {fmt_num(m['e2e_latency_std'], 3)}\n")
+            f.write(f"  Routing Lat (s):   {fmt_num(m['routing_latency_mean'], 4)} ± {fmt_num(m['routing_latency_std'], 4)}\n")
+            f.write(f"  Tool Lat (s):      {fmt_num(m['tool_latency_mean'], 3)} ± {fmt_num(m['tool_latency_std'], 3)}\n")
+            f.write("\n")
+    print(f"Generated plain-text summary: {txt_path}")
+
+
+if __name__ == "__main__":
+    main()

grpo_trainer.py

@@ -0,0 +1,840 @@
+import os
+import torch
+import torch.nn.functional as F
+import logging
+import csv
+from datetime import datetime
+from policy import RouterPolicy
+import copy
+import uuid
+
+# ---------------------------------------------------------------------------
+# Logger
+# ---------------------------------------------------------------------------
+logger = logging.getLogger("grpo_trainer")
+if not logger.handlers:
+    logger.setLevel(logging.INFO)
+    fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
+    ch = logging.StreamHandler()
+    ch.setFormatter(fmt)
+    logger.addHandler(ch)
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    logs_dir = os.path.join(repo_root, "logs")
+    try:
+        os.makedirs(logs_dir, exist_ok=True)
+    except Exception:
+        logs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    fh = logging.FileHandler(os.path.join(logs_dir, "grpo_trainer.log"))
+    fh.setFormatter(fmt)
+    logger.addHandler(fh)
+
+
+def _truncate(s: str, n: int = 500) -> str:
+    if s is None:
+        return ""
+    s = str(s)
+    return s if len(s) <= n else s[: n - 1] + "\u2026"
+
+
+# ---------------------------------------------------------------------------
+# RAG pipeline (unchanged from ppo_trainer)
+# ---------------------------------------------------------------------------
+def simulate_rag_pipeline(query, action_idx):
+    """
+    Maps action index to actual RAG modules via the backend API.
+    0: Unstructured / guidelines
+    1: Website / web-search
+    2: Structured / CBNU database
+    """
+    import sys
+
+    project_root = os.path.abspath(os.path.join(os.getcwd(), ".."))
+    backend_path = os.path.join(project_root, "backend")
+    sys.path.append(project_root)
+    sys.path.append(backend_path)
+
+    import asyncio
+    import requests
+
+    def detect_language(text):
+        if any("\uac00" <= char <= "\ud7a3" for char in text):
+            return "ko"
+        elif any("\u4e00" <= char <= "\u9fff" for char in text):
+            return "zh"
+        else:
+            return "en"
+
+    async def _call_tool(q, act_idx):
+        language = detect_language(q)
+        print(f"Detected language: {language}")
+        if act_idx == 0:
+            search_sources = ["guidelines"]
+        elif act_idx == 1:
+            search_sources = ["web-search"]
+        else:
+            search_sources = ["course-registration", "department-info"]
+
+        url = "http://localhost:7000/api/chat"
+        headers = {"accept": "application/json", "Content-Type": "application/json"}
+        payload = {
+            "message": q,
+            "language": language,
+            "search_sources": search_sources,
+            "session_id": f"grpo_train-{uuid.uuid4().hex}",
+        }
+        try:
+            response = requests.post(url, headers=headers, json=payload)
+            response.raise_for_status()
+            data = response.json()
+            answer = data.get("message", "")
+            context_parts = []
+            for src in data.get("sources", []):
+                if isinstance(src, dict) and src.get("success"):
+                    for res in src.get("data", {}).get("results", []):
+                        if "content" in res:
+                            context_parts.append(res["content"])
+                        elif "snippet" in res:
+                            context_parts.append(res["snippet"])
+            context = "\n".join(context_parts)
+            print("Tool Result fetched successfully via API.")
+            return context, answer
+        except Exception as e:
+            print(f"API call failed: {e}")
+            return "", ""
+
+    try:
+        ctx, ans = asyncio.run(_call_tool(query, int(action_idx)))
+        return ctx, ans
+    except Exception as e:
+        return "", f"simulate_rag_pipeline failed: {e}"
+
+
+# ---------------------------------------------------------------------------
+# LLM Judge Rewards (Gemini + OpenAI) — identical prompts to ppo_trainer
+# ---------------------------------------------------------------------------
+def llm_judge_reward(query, context, answer, action_idx):
+    import re
+    import requests
+
+    try:
+        model_name = "gemini-2.5-flash"
+        gemini_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent"
+        action_taken = f"Action {action_idx}"
+        gemini_key = "AIzaSyBoErDeUH1kRBbusI6qkMwISw5euyuXwZo"
+        gemini_token = os.getenv("GEMINI_API_TOKEN") or os.getenv("GCP_ACCESS_TOKEN")
+
+        Old_prompt = f"""You are an expert automated RAG evaluator.
+Your task is to grade the retrieved Context and the generated Answer based on the user's Query and the Action taken by the routing agent.
+
+### STRICT ROUTING RULES:
+The agent MUST select the correct action based on the query intent:
+- Action 0 (guidelines): Required for academic information, university rules, visa, scholarships, graduation requirements, health insurance, and part-time job info.
+- Action 1 (web_search): Required for university news, announcements, and external information.
+- Action 2 (cbnu_database): Required for department info (contacts, phone, email, website) and course registration (course details, schedules, professors, credits).
+
+### GRADING RUBRIC:
+1. Context Relevance (R_rel):
+- FATAL PENALTY: Did the agent choose an action that violates the STRICT ROUTING RULES above? If YES, you MUST score R_rel as 0.0 immediately, regardless of the context.
+- 0.0: Violates routing rules OR context is completely irrelevant/useless.
+- 0.5: Follows rules, but context contains only partial information.
+- 1.0: Follows rules AND context contains all necessary information to fully answer the query.
+
+2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
+- 0.0: Hallucinates information not in the context, or completely fails to answer the query.
+- 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
+- 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.
+
+### INPUTS:
+Query: {query}
+Action Taken: {action_taken}
+Context: {context}
+Answer: {answer}
+
+### OUTPUT FORMAT:
+You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
+Output this exact structure:
+{{
+"relevance_reasoning": "Brief 1-sentence explanation of why the rule was followed/broken and if context is relevant.",
+"R_rel": <float>,
+"correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
+"R_corr": <float>
+}}
+"""
+    
+        prompt = f"""You are an expert automated RAG evaluator.
+    Your task is to grade the retrieved Context and the generated Answer based on the user's Query.
+
+    ### GRADING RUBRIC:
+    1. Context Relevance (R_rel): How relevant is the retrieved Context to the user's Query?
+    - 0.0: The context is completely irrelevant or useless for answering the query.
+    - 0.5: The context is somewhat relevant but only contains partial information needed to answer the query.
+    - 1.0: The context is highly relevant and contains all necessary information to fully answer the query.
+
+    2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
+    - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
+    - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
+    - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.
+
+    ### INPUTS:
+    Query: {query}
+    Context: {context}
+    Answer: {answer}
+
+    ### OUTPUT FORMAT:
+    You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
+    Output this exact structure:
+    {{
+    "relevance_reasoning": "Brief 1-sentence explanation of why the context is or isn't relevant to the query.",
+    "R_rel": <float>,
+    "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
+    "R_corr": <float>
+    }}
+    """
+        headers = {"Content-Type": "application/json"}
+        params = {}
+        if gemini_key:
+            params = {"key": gemini_key}
+        elif gemini_token:
+            headers["Authorization"] = f"Bearer {gemini_token}"
+
+        payload = {
+            "contents": [{"parts": [{"text": prompt}]}],
+            "generationConfig": {"temperature": 0.0, "responseMimeType": "application/json"},
+        }
+        resp = requests.post(gemini_url, params=params, json=payload, headers=headers, timeout=15)
+        resp.raise_for_status()
+        data = resp.json()
+        try:
+            text = data["candidates"][0]["content"]["parts"][0]["text"]
+        except (KeyError, IndexError):
+            text = resp.text
+        m_rel = re.search(r'"R_rel"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        m_corr = re.search(r'"R_corr"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        if m_rel and m_corr:
+            R_rel, R_corr = float(m_rel.group(1)), float(m_corr.group(1))
+        else:
+            nums = re.findall(r"([0-9]*\.?[0-9]+)", text)
+            R_rel, R_corr = (float(nums[0]), float(nums[1])) if len(nums) >= 2 else (0.8, 0.9)
+    except Exception as e:
+        print(f"LLM judge failed, using fallback reward. Error: {e}", flush=True)
+        R_rel, R_corr = 0.7, 0.7
+
+    alpha, beta = 0.7, 0.3
+    print(f"LLM Judge for '{query[:50]}...' Scores => R_rel: {R_rel}, R_corr: {R_corr}", flush=True)
+    return float(R_rel * (alpha + beta * R_corr))
+
+
+def llm_judge_reward_openai(query, context, answer, action_idx):
+    import re
+    import requests
+
+    try:
+        model_name = "gpt-4o"
+        openai_url = "https://api.openai.com/v1/chat/completions"
+        action_taken = f"Action {action_idx}"
+        openai_key = "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"
+        if not openai_key:
+            print("OPENAI_API_KEY not found.", flush=True)
+            return 0.0
+
+        prompt = f"""You are an expert automated RAG evaluator.
+Your task is to grade the retrieved Context and the generated Answer based on the user's Query and the Action taken by the routing agent.
+
+### STRICT ROUTING RULES:
+The agent MUST select the correct action based on the query intent:
+- Action 0 (guidelines): Required for academic information, university rules, visa, scholarships, graduation requirements, health insurance, and part-time job info.
+- Action 1 (web_search): Required for university news, announcements, and external information.
+- Action 2 (cbnu_database): Required for department info (contacts, phone, email, website) and course registration (course details, schedules, professors, credits).
+
+### GRADING RUBRIC:
+1. Context Relevance (R_rel):
+- FATAL PENALTY: If the agent chose an action violating STRICT ROUTING RULES, score R_rel as 0.0 immediately.
+- 0.0: Violates routing rules OR context is completely irrelevant/useless.
+- 0.5: Follows rules, but context contains only partial information.
+- 1.0: Follows rules AND context contains all necessary information.
+
+2. Answer Correctness (R_corr): Is the Answer fully supported by the Context?
+- 0.0: Hallucinated or completely fails.
+- 0.5: Partially correct with minor issues.
+- 1.0: Perfectly accurate and faithful to context.
+
+### INPUTS:
+Query: {query}
+Action Taken: {action_taken}
+Context: {context}
+Answer: {answer}
+
+### OUTPUT FORMAT:
+Return ONLY a valid JSON object (no markdown blocks):
+{{
+"relevance_reasoning": "...",
+"R_rel": <float>,
+"correctness_reasoning": "...",
+"R_corr": <float>
+}}
+"""
+        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {openai_key}"}
+        payload = {
+            "model": model_name,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.0,
+            "response_format": {"type": "json_object"},
+        }
+        resp = requests.post(openai_url, json=payload, headers=headers, timeout=15)
+        resp.raise_for_status()
+        text = resp.json()["choices"][0]["message"]["content"]
+        m_rel = re.search(r'"R_rel"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        m_corr = re.search(r'"R_corr"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        if m_rel and m_corr:
+            R_rel, R_corr = float(m_rel.group(1)), float(m_corr.group(1))
+        else:
+            nums = re.findall(r"([0-9]*\.?[0-9]+)", text)
+            R_rel, R_corr = (float(nums[0]), float(nums[1])) if len(nums) >= 2 else (0.8, 0.9)
+    except Exception as e:
+        print(f"OpenAI LLM judge failed, using fallback reward. Error: {e}", flush=True)
+        R_rel, R_corr = 0.7, 0.7
+
+    alpha, beta = 0.7, 0.3
+    print(f"OpenAI Judge for '{query[:50]}...' Scores => R_rel: {R_rel}, R_corr: {R_corr}", flush=True)
+    return float(R_rel * (alpha + beta * R_corr))
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint helpers (unchanged interface)
+# ---------------------------------------------------------------------------
+def load_router_policy(checkpoint_path: str, embedding_dim=None, device: str = "cpu") -> RouterPolicy:
+    return RouterPolicy.load(checkpoint_path, embedding_dim=embedding_dim, device=device)
+
+
+def save_policy_checkpoint(policy: RouterPolicy, path: str, optimizer=None, meta=None):
+    return policy.save_checkpoint(path, optimizer=optimizer, meta=meta)
+
+
+# ---------------------------------------------------------------------------
+# GRPO Core Update
+# ---------------------------------------------------------------------------
+def grpo_update(
+    model,
+    ref_model,
+    optimizer,
+    states,            # Tensor [B, embed_dim]
+    actions,           # Tensor [B]
+    old_log_probs,     # Tensor [B]   — log-probs from the sampling policy
+    group_advantages,  # Tensor [B]   — already normalised per group
+    current_epoch: int,
+    total_epochs: int,
+    clip_ratio: float = 0.2,
+    kl_coef: float = 0.04,
+):
+    """
+    GRPO (Group Relative Policy Optimisation) update step.
+
+    Key differences vs PPO:
+    - No value/critic network: advantages are computed as within-group reward
+      normalisation (mean/std over the G samples drawn for the same query).
+    - KL divergence against a frozen reference policy is added as a regulariser
+      instead of (or alongside) the clipped surrogate.
+    - Dynamic entropy annealing is preserved from the original PPO implementation
+      as a paper novelty.
+    """
+    model.train()
+    device = next(model.parameters()).device
+
+    states        = states.to(device)
+    actions       = actions.to(device)
+    old_log_probs = old_log_probs.to(device)
+    group_advantages = group_advantages.to(device)
+
+    # ---- Current policy forward pass (actor logits only) ----
+    # RouterNetwork.forward() returns (logits, value); we use only logits.
+    logits, _ = model(states)
+
+    # Temperature scaling — keeps the same paper novelty as in ppo_update
+    temperature = 1.5
+    probs = F.softmax(logits / temperature, dim=-1)
+    dist  = torch.distributions.Categorical(probs)
+
+    new_log_probs = dist.log_prob(actions)
+    entropy       = dist.entropy().mean()
+
+    # ---- Clipped surrogate objective (same as PPO) ----
+    ratios = torch.exp(new_log_probs - old_log_probs)
+    surr1  = ratios * group_advantages
+    surr2  = torch.clamp(ratios, 1.0 - clip_ratio, 1.0 + clip_ratio) * group_advantages
+    actor_loss = -torch.min(surr1, surr2).mean()
+
+    # ---- KL penalty against frozen reference policy ----
+    with torch.no_grad():
+        ref_logits, _ = ref_model(states)
+        ref_probs = F.softmax(ref_logits / temperature, dim=-1)
+
+    # KL(current || reference) — per-sample, then averaged
+    kl_div = F.kl_div(
+        F.log_softmax(logits / temperature, dim=-1),
+        ref_probs,
+        reduction="batchmean",
+        log_target=False,
+    )
+
+    # ---- Dynamic Entropy Annealing (paper novelty, preserved) ----
+    # Start high (0.2) to break SFT overconfidence, decay to 0.01 to exploit.
+    initial_entropy = 0.30#0.2
+    final_entropy   = 0.02#0.01
+    decay_rate      = current_epoch / max(1, total_epochs - 1)
+    
+    ##Your entropy annealed from 0.20 → 0.01 over 20 epochs — that's too fast. With a better dataset and action 1 now viable, give the policy more room to explore early:
+    
+    dynamic_entropy_coef = initial_entropy - decay_rate * (initial_entropy - final_entropy)
+
+    # ---- Total GRPO loss ----
+    # L = actor_loss + kl_penalty - entropy_bonus
+    # (No critic/value loss — this is the core GRPO distinction from PPO)
+    total_loss = actor_loss + kl_coef * kl_div - dynamic_entropy_coef * entropy
+
+    optimizer.zero_grad()
+    total_loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
+    optimizer.step()
+
+    return total_loss.item(), dynamic_entropy_coef, kl_div.item()
+
+
+# ---------------------------------------------------------------------------
+# GRPO Rollout
+# ---------------------------------------------------------------------------
+def run_grpo_rollout_old(
+    model,
+    embedder,
+    queries,
+    optimizer,
+    epochs: int = 40,
+    group_size: int = 4,        # G — number of action samples per query
+    device: str = "cpu",
+    kl_coef: float = 0.04,
+    clip_ratio: float = 0.2,
+):
+    """
+    GRPO training loop for the RAG routing agent.
+
+    For each query we sample `group_size` (G) actions from the current policy,
+    collect a reward for each, then compute the within-group normalised advantage:
+
+        A_i = (r_i - mean(r_1..r_G)) / (std(r_1..r_G) + eps)
+
+    This replaces the PPO critic entirely — no value network is needed.
+
+    A frozen reference policy (snapshot of the model at the start of each epoch)
+    is used to compute the KL regularisation term in grpo_update().
+    """
+    import copy
+
+    model.train()
+
+    # ---- CSV debug log ----
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    csv_path  = os.path.join(repo_root, "logs", "grpo_debug.csv")
+    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
+    write_header = not os.path.exists(csv_path)
+    csv_file   = open(csv_path, mode="a", newline="", encoding="utf-8")
+    csv_writer = csv.writer(csv_file)
+    if write_header:
+        csv_writer.writerow([
+            "timestamp", "epoch", "step", "query",
+            "sample_idx", "action", "action_log_prob", "reward",
+            "group_advantage", "context", "answer",
+            "epoch_avg_reward", "grpo_loss", "kl_div",
+        ])
+
+    for epoch in range(epochs):
+
+        # ---- Freeze reference policy at the start of each epoch ----
+        ref_model = copy.deepcopy(model).to(device)
+        ref_model.eval()
+        for p in ref_model.parameters():
+            p.requires_grad_(False)
+
+        # Accumulators for the full batch (all queries × group_size)
+        all_states        = []
+        all_actions       = []
+        all_log_probs     = []
+        all_advantages    = []
+        epoch_rewards     = []
+
+        # ---- Rollout: collect G samples per query ----
+        for step_idx, query in enumerate(queries):
+            with torch.no_grad():
+                state_emb = embedder.encode(query, convert_to_tensor=True)
+                try:
+                    state_emb = state_emb.clone().to(device)
+                except Exception:
+                    state_emb = state_emb.clone().to(next(model.parameters()).device)
+
+            group_actions    = []
+            group_log_probs  = []
+            group_rewards    = []
+            group_contexts   = []
+            group_answers    = []
+
+            for g in range(group_size):
+                with torch.no_grad():
+                    # get_action returns (action, log_prob, value); ignore value in GRPO
+                    action, log_prob, _value = model.get_action(state_emb)
+
+                context, answer = simulate_rag_pipeline(query, action.item())
+                reward = llm_judge_reward(query, context, answer, action.item())
+
+                group_actions.append(action)
+                group_log_probs.append(log_prob)
+                group_rewards.append(reward)
+                group_contexts.append(context)
+                group_answers.append(answer)
+
+                try:
+                    lp_val = float(log_prob.item())
+                except Exception:
+                    lp_val = None
+
+                logger.info(
+                    "Epoch %d | Step %d | Sample %d | action=%d lp=%s reward=%.4f",
+                    epoch, step_idx, g, int(action.item()), lp_val, float(reward),
+                )
+
+            # ---- Within-group advantage normalisation ----
+            rewards_t = torch.tensor(group_rewards, dtype=torch.float32)
+            if group_size > 1:
+                adv = (rewards_t - rewards_t.mean()) / (rewards_t.std(unbiased=False) + 1e-8)
+            else:
+                adv = rewards_t - rewards_t.mean()
+
+            epoch_rewards.extend(group_rewards)
+
+            for g in range(group_size):
+                all_states.append(state_emb)
+                all_actions.append(group_actions[g])
+                all_log_probs.append(group_log_probs[g])
+                all_advantages.append(adv[g])
+
+                try:
+                    lp_val = float(group_log_probs[g].item())
+                except Exception:
+                    lp_val = None
+
+                csv_writer.writerow([
+                    datetime.utcnow().isoformat(), epoch, step_idx,
+                    _truncate(query, 800), g, int(group_actions[g].item()),
+                    lp_val, float(group_rewards[g]), float(adv[g]),
+                    _truncate(group_contexts[g], 1000), _truncate(group_answers[g], 1000),
+                    "", "", "",
+                ])
+
+        # ---- Build tensors for the optimisation phase ----
+        states_tensor     = torch.stack(all_states).detach()
+        actions_tensor    = torch.stack(all_actions).detach()
+        log_probs_tensor  = torch.stack(all_log_probs).detach()
+        advantages_tensor = torch.tensor(
+            [a.item() if hasattr(a, "item") else float(a) for a in all_advantages],
+            dtype=torch.float32,
+        ).to(device).detach()
+
+        # Global normalisation across all groups in the epoch (optional, improves stability)
+        if advantages_tensor.numel() > 1:
+            advantages_tensor = (
+                (advantages_tensor - advantages_tensor.mean())
+                / (advantages_tensor.std(unbiased=False) + 1e-8)
+            )
+
+        # ---- GRPO inner optimisation loop ----
+        grpo_inner_epochs = 4
+        epoch_loss = 0.0
+        last_kl    = 0.0
+
+        for _ in range(grpo_inner_epochs):
+            loss, ent_coef, kl = grpo_update(
+                model,
+                ref_model,
+                optimizer,
+                states_tensor,
+                actions_tensor,
+                log_probs_tensor,
+                advantages_tensor,
+                current_epoch=epoch,
+                total_epochs=epochs,
+                clip_ratio=clip_ratio,
+                kl_coef=kl_coef,
+            )
+            epoch_loss += loss
+            last_kl     = kl
+
+        avg_reward = float(sum(epoch_rewards) / len(epoch_rewards)) if epoch_rewards else 0.0
+        avg_loss   = epoch_loss / grpo_inner_epochs
+
+        logger.info(
+            "GRPO Epoch %d | Avg Reward: %.4f | Loss: %.4f | KL: %.6f | Ent Coef: %.4f",
+            epoch + 1, avg_reward, avg_loss, last_kl, ent_coef,
+        )
+        csv_writer.writerow([
+            datetime.utcnow().isoformat(), epoch, "", "", "", "", "", "",
+            "", "", "", avg_reward, avg_loss, last_kl,
+        ])
+        csv_file.flush()
+
+        print(
+            f"Epoch {epoch+1}/{epochs} | Avg Reward: {avg_reward:.4f} | "
+            f"Loss: {avg_loss:.4f} | KL: {last_kl:.6f} | Ent Coef: {ent_coef:.4f}"
+        )
+
+    try:
+        csv_file.close()
+    except Exception:
+        pass
+
+NUM_ACTIONS = 3   # Action 0: guidelines | Action 1: web-search | Action 2: database
+# model,
+# ref_model,
+# optimizer,
+# states,            # Tensor [B, embed_dim]
+# actions,           # Tensor [B]
+# old_log_probs,     # Tensor [B]   — log-probs from the sampling policy
+# group_advantages,  # Tensor [B]   — already normalised per group
+# current_epoch: int,
+# total_epochs: int,
+# clip_ratio: float = 0.2,
+# kl_coef: float = 0.04,
+
+def run_grpo_rollout(
+    model,
+    embedder,
+    queries,
+    optimizer,
+    ref_policy=None,      # ← added ref_policy argument to avoid re-creating inside trainer
+    epochs: int = 40,
+    start_epoch: int = 0,
+    device: str = "cpu",
+    kl_coef: float = 0.04,
+    clip_ratio: float = 0.2,
+):
+    """
+    GRPO training loop using EXHAUSTIVE group evaluation.
+
+    Why exhaustive instead of stochastic sampling?
+    -----------------------------------------------
+    Because the reward function is DETERMINISTIC:
+      same query + same action  →  same RAG output  →  same LLM-judge score.
+
+    Stochastic GRPO would keep drawing the same action from a confident policy,
+    producing std ≈ 0 → zero advantage → dead gradients.
+
+    Fix: for EVERY query, evaluate ALL 3 tools explicitly and use the reward
+    triplet [r0, r1, r2] as the group. Guaranteed non-zero variance as long
+    as any two tools score differently (almost always true for routing).
+
+    Group advantage formula (identical math to standard GRPO):
+        A_i = (r_i - mean(r0, r1, r2)) / (std(r0, r1, r2) + eps)
+
+    API calls per epoch: len(queries) × 3
+    """
+    model.train()
+
+    # ---- CSV debug log --------------------------------------------------- #
+    repo_root  = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+    csv_path   = os.path.join(repo_root, "logs", "grpo_debug.csv")
+    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
+    write_header = not os.path.exists(csv_path)
+    csv_file   = open(csv_path, mode="a", newline="", encoding="utf-8")
+    csv_writer = csv.writer(csv_file)
+    if write_header:
+        csv_writer.writerow([
+            "timestamp", "epoch", "step", "query",
+            "action", "reward", "group_advantage",
+            "context", "answer",
+            "epoch_avg_reward", "grpo_loss", "kl_div",
+        ])
+
+    for epoch in range(start_epoch, epochs):
+        # ── Reference policy setup ──────────────────────────────────────
+        if ref_policy is not None:
+            ref_model = ref_policy
+        elif epoch > 0 and epoch % 5 == 0:
+            # Refresh reference every 5 epochs to prevent KL explosion
+            ref_model = copy.deepcopy(model).to(device)
+            ref_model.eval()
+            for p in ref_model.parameters():
+                p.requires_grad_(False)
+            logger.info(f"Reference policy refreshed at epoch {epoch}")
+        else:
+            # First epoch OR non-refresh epoch — keep existing ref_model from previous epoch
+            if epoch == 0:
+                ref_model = copy.deepcopy(model).to(device)
+                ref_model.eval()
+                for p in ref_model.parameters():
+                    p.requires_grad_(False)
+            #else: ref_model carries over from previous epoch, no action needed
+
+        # Batch accumulators  (B = len(queries) × NUM_ACTIONS)
+        all_states     = []
+        all_actions    = []
+        all_log_probs  = []
+        all_advantages = []
+        epoch_rewards  = []
+
+        # ---- Rollout: evaluate ALL 3 actions for every query ------------- #
+        for step_idx, query in enumerate(queries):
+
+            # Embed the query ONCE — shared across all 3 action evaluations
+            with torch.no_grad():
+                state_emb = embedder.encode(query, convert_to_tensor=True)
+                try:
+                    state_emb = state_emb.clone().to(device)
+                except Exception:
+                    state_emb = state_emb.clone().to(next(model.parameters()).device)
+
+            group_rewards  = []
+            group_contexts = []
+            group_answers  = []
+
+            # Step 1: collect rewards for ALL 3 actions deterministically
+            for action_idx in range(NUM_ACTIONS):
+                context, answer = simulate_rag_pipeline(query, action_idx)
+                reward = llm_judge_reward(query, context, answer, action_idx)
+                #remove this hack once web-search tool is improved to not always score zero
+                if action_idx == 1 and reward == 0.0:
+                    # Don't let web-search always score zero — check if at least the routing was correct
+                    # Give a small non-zero signal to keep action 1 alive
+                    reward = 0.05   # tiny floor, just enough to prevent full collapse
+                group_rewards.append(reward)
+                group_contexts.append(context)
+                group_answers.append(answer)
+
+                logger.info(
+                    "Epoch %d | Step %d | action=%d reward=%.4f",
+                    epoch, step_idx, action_idx, float(reward),
+                )
+
+            # Step 2: within-group advantage normalisation
+            rewards_t = torch.tensor(group_rewards, dtype=torch.float32)
+            if rewards_t.std(unbiased=False) > 1e-8:
+                # Normal case: tools scored differently → strong learning signal
+                adv = (rewards_t - rewards_t.mean()) / (rewards_t.std(unbiased=False) + 1e-8)
+            else:
+                # Edge case: all 3 tools scored identically → zero advantage.
+                # Correct behaviour — do not push the policy in any direction.
+                adv = torch.zeros_like(rewards_t)
+                logger.warning(
+                    "Epoch %d | Step %d | All 3 actions equal reward=%.4f — zero gradient for this query.",
+                    epoch, step_idx, float(rewards_t[0]),
+                )
+
+            epoch_rewards.extend(group_rewards)
+
+            # Step 3: read log-probs for all 3 actions from the current policy
+            # in a single forward pass (no sampling needed)
+            with torch.no_grad():
+                logits, _ = model(state_emb.unsqueeze(0))
+                temperature = 1.5
+                probs = F.softmax(logits / temperature, dim=-1).squeeze(0)
+                log_probs_all = torch.log(probs + 1e-10)  # shape: [NUM_ACTIONS]
+
+            for action_idx in range(NUM_ACTIONS):
+                all_states.append(state_emb)
+                all_actions.append(torch.tensor(action_idx, dtype=torch.long))
+                all_log_probs.append(log_probs_all[action_idx])
+                all_advantages.append(adv[action_idx])
+
+                try:
+                    lp_val = float(log_probs_all[action_idx].item())
+                except Exception:
+                    lp_val = None
+
+                csv_writer.writerow([
+                    datetime.utcnow().isoformat(), epoch, step_idx,
+                    _truncate(query, 800),
+                    action_idx, float(group_rewards[action_idx]),
+                    float(adv[action_idx]),
+                    _truncate(group_contexts[action_idx], 1000),
+                    _truncate(group_answers[action_idx], 1000),
+                    "", "", "",
+                ])
+
+        # ---- Build tensors for the optimisation phase -------------------- #
+        states_tensor     = torch.stack(all_states).detach()
+        actions_tensor    = torch.stack(all_actions).detach()
+        log_probs_tensor  = torch.stack(all_log_probs).detach()
+        advantages_tensor = torch.tensor(
+            [a.item() if hasattr(a, "item") else float(a) for a in all_advantages],
+            dtype=torch.float32,
+        ).to(device).detach()
+
+        # Global re-normalisation across all (query × action) pairs in the epoch
+        if advantages_tensor.numel() > 1 and advantages_tensor.std(unbiased=False) > 1e-8:
+            advantages_tensor = (
+                (advantages_tensor - advantages_tensor.mean())
+                / (advantages_tensor.std(unbiased=False) + 1e-8)
+            )
+
+        # ---- GRPO inner optimisation loop (4 passes over same rollout) --- #
+        grpo_inner_epochs = 4
+        epoch_loss    = 0.0
+        last_kl       = 0.0
+        last_ent_coef = 0.0
+
+        for _ in range(grpo_inner_epochs):
+            loss, ent_coef, kl = grpo_update(
+                model,
+                ref_model,
+                optimizer,
+                states_tensor,
+                actions_tensor,
+                log_probs_tensor,
+                advantages_tensor,
+                current_epoch=epoch,
+                total_epochs=epochs,
+                clip_ratio=clip_ratio,
+                kl_coef=kl_coef,
+            )
+            epoch_loss   += loss
+            last_kl       = kl
+            last_ent_coef = ent_coef
+
+        avg_reward = float(sum(epoch_rewards) / len(epoch_rewards)) if epoch_rewards else 0.0
+        avg_loss   = epoch_loss / grpo_inner_epochs
+
+        logger.info(
+            "GRPO Epoch %d | Avg Reward: %.4f | Loss: %.4f | KL: %.6f | Ent Coef: %.4f",
+            epoch + 1, avg_reward, avg_loss, last_kl, last_ent_coef,
+        )
+        csv_writer.writerow([
+            datetime.utcnow().isoformat(), epoch, "", "", "", "", "", "", "",
+            avg_reward, avg_loss, last_kl,
+        ])
+        csv_file.flush()
+
+        print(
+            f"Epoch {epoch+1}/{epochs} | "
+            f"Avg Reward: {avg_reward:.4f} | "
+            f"Loss: {avg_loss:.4f} | "
+            f"KL: {last_kl:.6f} | "
+            f"Ent Coef: {last_ent_coef:.4f}"
+        )
+        # ---- Save checkpoint after every epoch --------------------------- #
+        adapters_dir = os.path.join(os.path.dirname(__file__), "adapters")
+        os.makedirs(adapters_dir, exist_ok=True)
+        epoch_ckpt_path = os.path.join(adapters_dir, f"router_grpo_epoch_{epoch+1}.pt")
+        torch.save(
+            {
+                "epoch": epoch + 1,
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "avg_reward": avg_reward,
+                "avg_loss": avg_loss,
+                "kl_div": last_kl,
+                "ent_coef": last_ent_coef,
+            },
+            epoch_ckpt_path,
+        )
+        logger.info("Saved checkpoint → %s", epoch_ckpt_path)
+        print(f"  └─ Checkpoint saved: adapters/router_grpo_epoch_{epoch+1}.pt")
+
+    try:
+        csv_file.close()
+    except Exception:
+        pass

main.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+import os
+import sys
+import argparse
+import csv
+import torch
+
+ROOT = os.path.dirname(__file__)
+
+def build_dataset_from_csv(csv_path):
+    dataset = []
+    with open(csv_path, newline='', encoding='utf-8') as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            q = row.get('question') or row.get('Question')
+            lbl = row.get('label') or row.get('Label')
+            if q is None or lbl is None:
+                continue
+            try:
+                lbl = int(lbl)
+            except Exception:
+                continue
+            dataset.append({'question': q, 'label': lbl})
+    return dataset
+
+
+class EmbedderWrapper:
+    """Wrap a SentenceTransformer instance to ensure tensors are on `device`.
+
+    The wrapper exposes `encode(texts, convert_to_tensor=True, **kwargs)` same as
+    the SentenceTransformer API and will move returned torch.Tensor to the device.
+    """
+    def __init__(self, embedder, device):
+        self.embedder = embedder
+        self.device = device
+
+    def encode(self, texts, convert_to_tensor=True, **kwargs):
+        emb = self.embedder.encode(texts, convert_to_tensor=convert_to_tensor, **kwargs)
+        # Move tensor outputs to device (if it's a torch.Tensor)
+        if convert_to_tensor and hasattr(emb, 'to'):
+            try:
+                emb = emb.to(self.device)
+            except Exception:
+                pass
+        return emb
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Run SFT training for RouterNetwork')
+    parser.add_argument('--csv', default='processed_dataset.csv', help='Path to CSV dataset')
+    parser.add_argument('--model_name', default='all-MiniLM-L6-v2', help='SentenceTransformer model name')
+    parser.add_argument('--epochs', type=int, default=5)
+    parser.add_argument('--lr', type=float, default=1e-3)
+    parser.add_argument('--checkpoint', default='router_checkpoint.pt', help='Checkpoint output path')
+    parser.add_argument('--device', default=None, help='Device to use (cpu or cuda). Auto-detect if unset')
+    args = parser.parse_args()
+
+    # Ensure routing_agent directory is importable such that its modules' imports work
+    routing_agent_path = os.path.join(ROOT, 'routing_agent')
+    if routing_agent_path not in sys.path:
+        sys.path.insert(0, routing_agent_path)
+
+    # Import modules from routing_agent (sft_trainer expects to import model as a top-level module)
+    try:
+        import sft_trainer
+        import model as model_mod
+    except Exception as e:
+        print('Failed importing routing_agent modules:', e)
+        sys.exit(1)
+
+    # Device selection
+    if args.device:
+        device = torch.device(args.device)
+    else:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Build dataset
+    csv_path = os.path.join(ROOT, args.csv) if not os.path.isabs(args.csv) else args.csv
+    if not os.path.exists(csv_path):
+        print('Dataset CSV not found at', csv_path)
+        sys.exit(1)
+
+    dataset = build_dataset_from_csv(csv_path)
+    if len(dataset) == 0:
+        print('No valid rows found in dataset. Exiting.')
+        sys.exit(1)
+
+    # Initialize embedder
+    try:
+        from sentence_transformers import SentenceTransformer
+    except Exception as e:
+        print('sentence-transformers not installed or failed to import:', e)
+        sys.exit(1)
+
+    embedder = SentenceTransformer(args.model_name)
+    embedder_wrapper = EmbedderWrapper(embedder, device)
+
+    # Infer embedding dim and instantiate model from the same module sft_trainer expects
+    try:
+        embed_dim = embedder.get_sentence_embedding_dimension()
+    except Exception:
+        # fallback to common dim
+        embed_dim = 384
+
+    model = model_mod.RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)
+
+    # Run SFT training
+    print(f"Starting SFT training on device={device} with model={args.model_name} | epochs={args.epochs} | lr={args.lr}")
+    sft_trainer.train_sft(model, embedder_wrapper, dataset, epochs=args.epochs, lr=args.lr)
+
+    # Save checkpoint
+    try:
+        torch.save(model.state_dict(), args.checkpoint)
+        print('Saved checkpoint to', args.checkpoint)
+    except Exception as e:
+        print('Failed to save checkpoint:', e)
+
+
+if __name__ == '__main__':
+    main()

model.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sentence_transformers import SentenceTransformer
+
+class RouterNetwork(nn.Module):
+    def __init__(self, embedding_dim = 384, hidden_dim = 128, num_classes = 3):
+        super(RouterNetwork, self).__init__()
+        
+        # shared representation layer
+        self.shared_layer = nn.Sequential(
+            nn.Linear(embedding_dim, hidden_dim),
+            nn.ReLU(),
+            nn.LayerNorm(hidden_dim)
+        )
+        
+        # Actor Head: outputs logits for the 3 modules
+        # 0: unstructured, 1: website, 2: structured
+        self.actor_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_dim // 2, num_classes)
+        )
+        
+        # Critic Head: outputs a single value for the state value estimation
+        self.critic_head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(),
+            nn.Linear(hidden_dim // 2, 1)
+        )
+        
+    def forward(self, x):
+        shared_features = self.shared_layer(x)
+        action_logits = self.actor_head(shared_features)
+        state_value = self.critic_head(shared_features)
+        return action_logits, state_value
+    
+    def get_action(self, x):
+        action_logits, value = self.forward(x)
+        action_probs = F.softmax(action_logits, dim=-1)
+        dist = torch.distributions.Categorical(action_probs)
+        action = dist.sample()
+        return action, dist.log_prob(action), value

policy.py

@@ -0,0 +1,83 @@
+import os
+import torch
+import logging
+
+logger = logging.getLogger(__name__)
+
+from model import RouterNetwork
+
+
+class RouterPolicy:
+    """Thin wrapper around RouterNetwork for loading, inference and checkpointing."""
+
+    def __init__(self, model: RouterNetwork, device: str = "cpu"):
+        self.model = model.to(device)
+        self.device = torch.device(device)
+        self.model.eval()
+
+    @classmethod
+    def load(cls, checkpoint_path: str, embedding_dim: int | None = None, device: str = "cpu"):
+        """Load RouterNetwork from checkpoint. If checkpoint contains container dict with
+        'model_state_dict' will use that, otherwise will try to load state_dict directly.
+
+        embedding_dim may be required if checkpoint does not include meta info.
+        """
+        checkpoint_path = os.path.abspath(checkpoint_path)
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"Router checkpoint not found: {checkpoint_path}")
+
+        sd = torch.load(checkpoint_path, map_location="cpu")
+
+        # Unpack common formats
+        if isinstance(sd, dict) and "model_state_dict" in sd:
+            model_state = sd["model_state_dict"]
+            meta = sd.get("meta", {})
+        else:
+            model_state = sd
+            meta = {}
+
+        # Infer embedding dim from meta if available
+        emb_dim = embedding_dim or meta.get("embedding_dim") or meta.get("embed_dim") or 384
+
+        model = RouterNetwork(embedding_dim=emb_dim)
+
+        missing, unexpected = model.load_state_dict(model_state, strict=False)
+        if missing or unexpected:
+            logger.info(f"RouterPolicy.load: missing keys={list(missing)} unexpected={list(unexpected)}")
+
+        return cls(model, device=device)
+
+    def predict_sync(self, emb_tensor: torch.Tensor):
+        """Synchronous prediction: accepts a 1-D or 2-D embedding tensor and returns
+        (action_idx:int, prob:float, value:float).
+        """
+        was_1d = False
+        if emb_tensor.dim() == 1:
+            emb_tensor = emb_tensor.unsqueeze(0)
+            was_1d = True
+
+        emb = emb_tensor.to(self.device)
+
+        with torch.no_grad():
+            logits, value = self.model(emb)
+            probs = torch.softmax(logits, dim=-1)
+            top_idx = int(probs.argmax(dim=-1)[0].item())
+            top_prob = float(probs[0, top_idx].item())
+            val = float(value.squeeze(1)[0].item()) if value is not None else 0.0
+
+        return top_idx, top_prob, val
+
+    def save_checkpoint(self, path: str, optimizer: torch.optim.Optimizer | None = None, meta: dict | None = None):
+        path = os.path.abspath(path)
+        payload = {
+            "model_state_dict": self.model.state_dict(),
+            "meta": meta or {}
+        }
+        if optimizer is not None:
+            try:
+                payload["optimizer_state_dict"] = optimizer.state_dict()
+            except Exception:
+                logger.warning("Could not include optimizer state in checkpoint")
+
+        torch.save(payload, path)
+        return path

ppo_trainer.py

@@ -0,0 +1,772 @@
+import os
+import torch
+import torch.optim as optim
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+import csv
+from datetime import datetime
+from policy import RouterPolicy
+
+
+# Initialize module logger
+logger = logging.getLogger("ppo_trainer")
+if not logger.handlers:
+    logger.setLevel(logging.INFO)
+    fmt = logging.Formatter("%(asctime)s %(levelname)s %(message)s")
+    # Console handler
+    ch = logging.StreamHandler()
+    ch.setFormatter(fmt)
+    logger.addHandler(ch)
+    # File handler (append) - place logs directory at repository root `logs/`
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+    logs_dir = os.path.join(repo_root, 'logs')
+    try:
+        os.makedirs(logs_dir, exist_ok=True)
+    except Exception:
+        logs_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+
+    fh_path = os.path.join(logs_dir, 'ppo_trainer.log')
+    fh = logging.FileHandler(fh_path)
+    fh.setFormatter(fmt)
+    logger.addHandler(fh)
+
+
+def _truncate(s: str, n: int = 500) -> str:
+    if s is None:
+        return ""
+    s = str(s)
+    return s if len(s) <= n else s[:n-1] + "\u2026"
+
+
+def v1_ppo_update(model, optimizer, states, actions, old_log_probs, returns, advantages, clip_ratio=0.2):
+    """
+    Performs a single PPO update step.
+    
+    states: Tensor of query embeddings
+    actions: Tensor of actions taken
+    returns: Tensor of discounted cumulative rewards
+    advantages: Tensor of advantages (Returns - State Values)
+    """
+    model.train()
+    device = next(model.parameters()).device
+
+    # ensure tensors live on the same device as the model
+    states = states.to(device)
+    actions = actions.to(device)
+    old_log_probs = old_log_probs.to(device)
+    returns = returns.to(device)
+    advantages = advantages.to(device)
+
+    #1. Get current policy probabilities and values
+    logits, current_values = model(states)
+    probs = F.softmax(logits, dim=-1)
+    dist = torch.distributions.Categorical(probs)
+    
+    # calculate new log probabilities and entropy for the taken actions
+    new_log_probs = dist.log_prob(actions)
+    entropy = dist.entropy().mean()
+    
+    #2. calculate the ratio: pi_theta(a|s) / pi_theta_old(a|s)
+    ratios = torch.exp(new_log_probs - old_log_probs)
+    
+    #3. calculate surrogate losses
+    surr1 = ratios * advantages
+    surr2 = torch.clamp(ratios, 1.0 - clip_ratio, 1.0 + clip_ratio) * advantages
+    
+    # calculate actor and critic losses
+    # PPO loss is the negative of the minimum surrogate loss plus value function loss
+    actor_loss = -torch.min(surr1, surr2).mean()  # PPO objective (maximize the minimum of the surrogate losses)
+    critic_loss = F.mse_loss(current_values.squeeze(), returns.view_as(current_values.squeeze()))  # Value function loss
+    
+    #5. total Loss calculation (with entropy bonus to encourage exploration)
+    entropy_coef = 0.08 #increased after v1 experiment #0.01 # common weight for the entropy bonus
+    value_coef = 0.5 # common weight for the value loss
+    total_loss = actor_loss + (value_coef * critic_loss) - (entropy_coef * entropy)  # Combine losses with appropriate weights
+    
+    #6. backpropagation and optimization step
+    optimizer.zero_grad()
+    total_loss.backward()
+    
+    # Gradient clipping (optional but common in PPO to stabilize training)
+    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
+    optimizer.step()
+    
+    return total_loss.item()
+    
+def ppo_update(model, optimizer, states, actions, old_log_probs, returns, advantages, current_epoch, total_epochs, clip_ratio=0.2):
+    model.train()
+    
+    # 1. Forward pass
+    logits, current_values = model(states)
+    
+    # Optional: Logit Temperature Scaling (helps soften the SFT spike)
+    temperature = 1.5 
+    probs = F.softmax(logits / temperature, dim=-1)
+    dist = torch.distributions.Categorical(probs)
+    
+    new_log_probs = dist.log_prob(actions)
+    entropy = dist.entropy().mean()
+    
+    # 2. Calculate PPO Ratio
+    ratios = torch.exp(new_log_probs - old_log_probs)
+    
+    # 3. Calculate Surrogate Losses
+    surr1 = ratios * advantages
+    surr2 = torch.clamp(ratios, 1.0 - clip_ratio, 1.0 + clip_ratio) * advantages
+    
+    actor_loss = -torch.min(surr1, surr2).mean()
+    critic_loss = F.mse_loss(current_values.squeeze(), returns)
+    
+    # ---------------------------------------------------------
+    # NOVELTY FOR THE PAPER: Dynamic Entropy Annealing
+    # Start high (0.2) to shatter SFT confidence, decay to low (0.01) to exploit
+    # ---------------------------------------------------------
+    initial_entropy = 0.2
+    final_entropy = 0.01
+    decay_rate = current_epoch / max(1, total_epochs - 1)
+    dynamic_entropy_coef = initial_entropy - decay_rate * (initial_entropy - final_entropy)
+    
+    value_coef = 0.5
+    total_loss = actor_loss + (value_coef * critic_loss) - (dynamic_entropy_coef * entropy)
+    
+    optimizer.zero_grad()
+    total_loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
+    optimizer.step()
+    
+    return total_loss.item(), dynamic_entropy_coef
+    
+def simulate_rag_pipeline(query, action_idx):
+    """
+    Dummy function mapping to your actual RAG modules.
+    0: Unstructured, 1: Website, 2: Structured (SQL)
+    """
+    # Integrate with backend SimpleAgent to run the actual tools.
+    # This function will synchronously run async agent methods using asyncio.run
+    import sys
+    import os
+    import asyncio
+    import requests
+
+    # Ensure repo root is on sys.path so we can import backend package
+    # repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+    # if repo_root not in sys.path:
+    #     sys.path.insert(0, repo_root)
+
+    import sys
+    import os
+
+    project_root = os.path.abspath(os.path.join(os.getcwd(), ".."))
+    backend_path = os.path.join(project_root, "backend")
+
+    sys.path.append(project_root)
+    sys.path.append(backend_path)
+
+    try:
+        from backend.app.core.agent import SimpleAgent
+    except Exception:
+        # Fallback import path for different execution contexts
+        try:
+            from app.core.agent import SimpleAgent
+        except Exception as e:
+            return "", f"Failed importing SimpleAgent: {e}"
+    def detect_language(text):
+        #check if any character in the text is in the Korean Unicode range or chinese Unicode range or english Unicode range
+        if any('\uac00' <= char <= '\ud7a3' for char in text):
+            return "ko"
+        elif any('\u4e00' <= char <= '\u9fff' for char in text):
+            return "zh"
+        else:
+            return "en"
+    
+    async def old_call_tool(q, act_idx):
+        agent = SimpleAgent()
+        language = detect_language(q)
+        if act_idx == 0:
+            res = await agent._execute_guidelines(q, language)
+        elif act_idx == 1:
+            res = await agent._execute_web_search(q, language)
+        else:
+            res = await agent._execute_cbnu_database(q, language)
+
+        # Build a short context from tool results
+        context = ""
+        try:
+            results = res.data.get("results", []) if res and res.data else []
+            snippets = []
+            for item in results[:3]:
+                snippet = item.get("snippet") or item.get("content") or item.get(f"answer_{language}") or str(item)
+                snippets.append(snippet)
+            context = "\n".join(snippets)
+        except Exception:
+            context = ""
+
+        # Use agent's generator to synthesize an answer from the single tool result
+        try:
+            answer = await agent._generate_response(q, [res], language)
+        except Exception:
+            answer = ""
+
+        return context, answer
+
+    async def _call_tool(q, act_idx):
+        language = detect_language(q)
+        print(f"Detected language: {language}")
+        
+        # 0: guidelines, 1: web search, 2: database
+        if act_idx == 0:
+            search_sources = ["guidelines"]
+        elif act_idx == 1:
+            search_sources = ["web-search"]
+        else:
+            search_sources = ["course-registration", "department-info"]
+            
+        url = 'http://localhost:7000/api/chat'
+        headers = {
+            'accept': 'application/json',
+            'Content-Type': 'application/json'
+        }
+        payload = {
+            "message": q,
+            "language": language,
+            "search_sources": search_sources,
+            "session_id": "ppo_train_session"
+        }
+
+        try:
+            response = requests.post(url, headers=headers, json=payload)
+            response.raise_for_status()
+            data = response.json()
+            
+            # Extract context and answer from the API response
+            answer = data.get("message", "")
+            
+            # Extract content securely from the sources
+            context_parts = []
+            sources = data.get("sources", [])
+            for src in sources:
+                if isinstance(src, dict) and src.get("success"):
+                    src_data = src.get("data", {})
+                    results = src_data.get("results", [])
+                    for res in results:
+                        if "content" in res:
+                            context_parts.append(res["content"])
+                        elif "snippet" in res:
+                            context_parts.append(res["snippet"])
+                            
+            context = "\n".join(context_parts)
+                
+            print(f"Tool Result fetched successfully via API.")
+            return context, answer
+
+        except Exception as e:
+            print(f"API call failed: {e}")
+            return "", ""
+    try:
+        ctx, ans = asyncio.run(_call_tool(query, int(action_idx)))
+        return ctx, ans
+    except Exception as e:
+        return "", f"simulate_rag_pipeline failed: {e}"
+
+def llm_judge_reward(query, context, answer, action_idx):
+    import os
+    import requests
+    import re
+    
+    try:
+        # 1. FIXED: Correct API version (v1beta) and method (:generateContent)
+        # Note: If gemini-2.5-flash still throws 404, change to gemini-2.0-flash or gemini-1.5-flash
+        model_name = "gemini-2.5-flash" 
+        gemini_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model_name}:generateContent"
+        action_taken = f"Action {action_idx}"
+        gemini_key = "AIzaSyBoErDeUH1kRBbusI6qkMwISw5euyuXwZo"
+        gemini_token = os.getenv("GEMINI_API_TOKEN") or os.getenv("GCP_ACCESS_TOKEN")
+        
+        Old_prompt = f"""You are an expert automated RAG evaluator. 
+        Your task is to grade the retrieved Context and the generated Answer based on the user's Query.
+
+        ### GRADING RUBRIC:
+        1. Context Relevance (R_rel): Does the Context contain the necessary information to answer the Query?
+        - 0.0: Completely irrelevant or useless.
+        - 0.5: Contains partial information, but not enough for a complete answer.
+        - 1.0: Contains all necessary information to fully answer the query.
+
+        2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
+        - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
+        - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
+        - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.
+
+        ### INPUTS:
+        Query: {query}
+        Context: {context}
+        Answer: {answer}
+
+        ### OUTPUT FORMAT:
+        You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
+        Output this exact structure:
+        {{
+            "relevance_reasoning": "Brief 1-sentence explanation of why the context is/isn't relevant.",
+            "R_rel": <float between 0.0 and 1.0>,
+            "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
+            "R_corr": <float between 0.0 and 1.0>
+        }}
+        """
+
+        prompt = f"""You are an expert automated RAG evaluator. 
+            Your task is to grade the retrieved Context and the generated Answer based on the user's Query and the Action taken by the routing agent.
+
+            ### STRICT ROUTING RULES:
+            The agent MUST select the correct action based on the query intent:
+            - Action 0 (guidelines): Required for academic information, university rules, visa, scholarships, graduation requirements, health insurance, and part-time job info.
+            - Action 1 (web_search): Required for university news, announcements, and external information.
+            - Action 2 (cbnu_database): Required for department info (contacts, phone, email, website) and course registration (course details, schedules, professors, credits).
+
+            ### GRADING RUBRIC:
+            1. Context Relevance (R_rel): 
+            - FATAL PENALTY: Did the agent choose an action that violates the STRICT ROUTING RULES above? If YES, you MUST score R_rel as 0.0 immediately, regardless of the context.
+            - 0.0: Violates routing rules OR context is completely irrelevant/useless.
+            - 0.5: Follows rules, but context contains only partial information.
+            - 1.0: Follows rules AND context contains all necessary information to fully answer the query.
+
+            2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
+            - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
+            - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
+            - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.
+
+            ### INPUTS:
+            Query: {query}
+            Action Taken: {action_taken}
+            Context: {context}
+            Answer: {answer}
+
+            ### OUTPUT FORMAT:
+            You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
+            Output this exact structure:
+            {{
+                "relevance_reasoning": "Brief 1-sentence explanation of why the rule was followed/broken and if context is relevant.",
+                "R_rel": <float between 0.0 and 1.0>,
+                "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
+                "R_corr": <float between 0.0 and 1.0>
+            }}
+            """
+        
+        headers = {"Content-Type": "application/json"}
+        params = {}
+        
+        if gemini_key:
+            params = {"key": gemini_key}
+        elif gemini_token:
+            headers["Authorization"] = f"Bearer {gemini_token}"
+
+        # 2. FIXED: Correct Gemini payload structure
+        payload = {
+            "contents": [{
+                "parts": [{"text": prompt}]
+            }],
+            "generationConfig": {
+                # "maxOutputTokens": 200,
+                "temperature": 0.0, # Deterministic output for consistent rewards
+                "responseMimeType": "application/json" # 3. ADDED: Enforces JSON output
+            }
+        }
+
+        resp = requests.post(gemini_url, params=params, json=payload, headers=headers, timeout=15)
+        
+        if resp.status_code == 404:
+            print(f"Gemini endpoint returned 404. Check if '{model_name}' is available to your API key.", flush=True)
+            print(f"Tried URL: {gemini_url}", flush=True)
+            resp.raise_for_status()
+            
+        resp.raise_for_status()
+
+        # 4. FIXED: Parsing the correct Gemini response path
+        data = resp.json()
+        text = ""
+        try:
+            text = data["candidates"][0]["content"]["parts"][0]["text"]
+        except (KeyError, IndexError):
+            text = resp.text # Fallback to raw text if structure is unexpected
+
+        # Extract floats using your existing regex
+        m_rel = re.search(r'"R_rel"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        m_corr = re.search(r'"R_corr"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        
+        if m_rel and m_corr:
+            R_rel = float(m_rel.group(1))
+            R_corr = float(m_corr.group(1))
+        else:
+            nums = re.findall(r'([0-9]*\.?[0-9]+)', text)
+            if len(nums) >= 2:
+                R_rel = float(nums[0])
+                R_corr = float(nums[1])
+            else:
+                R_rel, R_corr = 0.8, 0.9
+                
+    except Exception as e:
+        print("LLM judge failed, using fallback reward. Error details:", flush=True)
+        print(e, flush=True)
+        R_rel, R_corr = 0.7, 0.7
+
+    # The Reward Formulation: R_total = R_rel * (alpha + beta * R_corr)
+    alpha, beta = 0.7, 0.3
+    print(f"LLM Judge for '{query[:50]}...' Scores => R_rel: {R_rel}, R_corr: {R_corr}", flush=True)
+    reward = R_rel * (alpha + beta * R_corr)
+
+    return float(reward)
+
+
+def llm_judge_reward_openai(query, context, answer, action_idx):
+    import os
+    import requests
+    import re
+    
+    try:
+        model_name = "gpt-4o" 
+        openai_url = "https://api.openai.com/v1/chat/completions"
+        action_taken = f"Action {action_idx}"
+        openai_key = "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"
+        if not openai_key:
+            print("OPENAI_API_KEY not found in environment.", flush=True)
+            return 0.0
+            
+        old_prompt = f"""You are an expert automated RAG evaluator. 
+        Your task is to grade the retrieved Context and the generated Answer based on the user's Query.
+
+        ### GRADING RUBRIC:
+        1. Context Relevance (R_rel): Does the Context contain the necessary information to answer the Query?
+        - 0.0: Completely irrelevant or useless.
+        - 0.5: Contains partial information, but not enough for a complete answer.
+        - 1.0: Contains all necessary information to fully answer the query.
+
+        2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
+        - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
+        - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
+        - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.
+
+        ### INPUTS:
+        Query: {query}
+        Context: {context}
+        Answer: {answer}
+
+        ### OUTPUT FORMAT:
+        You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
+        Output this exact structure:
+        {{
+            "relevance_reasoning": "Brief 1-sentence explanation of why the context is/isn't relevant.",
+            "R_rel": <float between 0.0 and 1.0>,
+            "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
+            "R_corr": <float between 0.0 and 1.0>
+        }}
+        """
+        prompt = f"""You are an expert automated RAG evaluator. 
+            Your task is to grade the retrieved Context and the generated Answer based on the user's Query and the Action taken by the routing agent.
+
+            ### STRICT ROUTING RULES:
+            The agent MUST select the correct action based on the query intent:
+            - Action 0 (guidelines): Required for academic information, university rules, visa, scholarships, graduation requirements, health insurance, and part-time job info.
+            - Action 1 (web_search): Required for university news, announcements, and external information.
+            - Action 2 (cbnu_database): Required for department info (contacts, phone, email, website) and course registration (course details, schedules, professors, credits).
+
+            ### GRADING RUBRIC:
+            1. Context Relevance (R_rel): 
+            - FATAL PENALTY: Did the agent choose an action that violates the STRICT ROUTING RULES above? If YES, you MUST score R_rel as 0.0 immediately, regardless of the context.
+            - 0.0: Violates routing rules OR context is completely irrelevant/useless.
+            - 0.5: Follows rules, but context contains only partial information.
+            - 1.0: Follows rules AND context contains all necessary information to fully answer the query.
+
+            2. Answer Correctness (R_corr): Is the Answer fully supported by the Context? (IGNORE your internal knowledge. Grade strictly on the provided context).
+            - 0.0: Hallucinates information not in the context, or completely fails to answer the query.
+            - 0.5: Partially correct, but misses key details from the context or includes slight hallucinations.
+            - 1.0: Perfectly accurate, directly answers the query, and is 100% faithful to the context.
+
+            ### INPUTS:
+            Query: {query}
+            Action Taken: {action_taken}
+            Context: {context}
+            Answer: {answer}
+
+            ### OUTPUT FORMAT:
+            You must return ONLY a valid JSON object. Do not use markdown blocks (```json).
+            Output this exact structure:
+            {{
+                "relevance_reasoning": "Brief 1-sentence explanation of why the rule was followed/broken and if context is relevant.",
+                "R_rel": <float between 0.0 and 1.0>,
+                "correctness_reasoning": "Brief 1-sentence explanation of answer correctness and faithfulness.",
+                "R_corr": <float between 0.0 and 1.0>
+            }}
+            """
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {openai_key}"
+        }
+        
+        payload = {
+            "model": model_name,
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.0,
+            "response_format": {"type": "json_object"}
+        }
+
+        resp = requests.post(openai_url, json=payload, headers=headers, timeout=15)
+        resp.raise_for_status()
+
+        data = resp.json()
+        text = data["choices"][0]["message"]["content"]
+
+        # Extract floats
+        m_rel = re.search(r'"R_rel"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        m_corr = re.search(r'"R_corr"\s*:\s*([0-9]*\.?[0-9]+)', text)
+        
+        if m_rel and m_corr:
+            R_rel = float(m_rel.group(1))
+            R_corr = float(m_corr.group(1))
+        else:
+            nums = re.findall(r'([0-9]*\.?[0-9]+)', text)
+            if len(nums) >= 2:
+                R_rel = float(nums[0])
+                R_corr = float(nums[1])
+            else:
+                R_rel, R_corr = 0.8, 0.9
+                
+    except Exception as e:
+        print("OpenAI LLM judge failed, using fallback reward. Error details:", flush=True)
+        print(e, flush=True)
+        R_rel, R_corr = 0.7, 0.7
+
+    alpha, beta = 0.7, 0.3
+    print(f"OpenAI Judge for '{query[:50]}...' Scores => R_rel: {R_rel}, R_corr: {R_corr}", flush=True)
+    reward = R_rel * (alpha + beta * R_corr)
+
+    return float(reward)
+
+
+
+
+def load_router_policy(checkpoint_path: str, embedding_dim: int | None = None, device: str = "cpu") -> RouterPolicy:
+    """Load a RouterPolicy wrapper from checkpoint path."""
+    return RouterPolicy.load(checkpoint_path, embedding_dim=embedding_dim, device=device)
+
+
+def save_policy_checkpoint(policy: RouterPolicy, path: str, optimizer: torch.optim.Optimizer | None = None, meta: dict | None = None):
+    return policy.save_checkpoint(path, optimizer=optimizer, meta=meta)
+
+def v1_run_ppo_rollout(model, embedder, queries, optimizer, epochs=5, device: str = "cpu"):
+    """
+    Executes the environment loop and updates the model.
+    """
+    model.train()
+    gamma = 0.99 # Discount factor
+    
+    # Prepare CSV debug log (repo-level `logs/ppo_debug.csv`)
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+    csv_path = os.path.abspath(os.path.join(repo_root, 'logs', 'ppo_debug.csv'))
+    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
+    write_header = not os.path.exists(csv_path)
+    csv_file = open(csv_path, mode='a', newline='', encoding='utf-8')
+    csv_writer = csv.writer(csv_file)
+    if write_header:
+        csv_writer.writerow([
+            'timestamp', 'epoch', 'step', 'query', 'action', 'action_log_prob', 'value', 'reward', 'context', 'answer', 'epoch_avg_reward', 'policy_loss'
+        ])
+
+    for epoch in range(epochs):
+        states, actions, log_probs, rewards, values = [], [], [], [], []
+        
+        # 1. Rollout Phase: Gather Experience
+        for step_idx, query in enumerate(queries):
+            # Create embedding (SentenceTransformer may use inference_mode internally).
+            state_emb = embedder.encode(query, convert_to_tensor=True)
+
+            # Clone the embedding to ensure it's a regular tensor (not an "inference" tensor)
+            # and move it to the model/device.
+            try:
+                state_emb = state_emb.clone().to(device)
+            except Exception:
+                # fallback: try model parameter device
+                state_emb = state_emb.clone().to(next(model.parameters()).device)
+
+            # Agent selects a module
+            action, log_prob, value = model.get_action(state_emb)
+            
+            # Execute RAG and get LLM Judge Score
+            context, answer = simulate_rag_pipeline(query, action.item())
+            # reward = llm_judge_reward(query, context, answer, action.item())
+            reward = llm_judge_reward_openai(query, context, answer, action.item())
+
+            # Log per-step details to logger and CSV for debugging
+            try:
+                log_prob_val = float(log_prob.item()) if hasattr(log_prob, 'item') else float(log_prob)
+            except Exception:
+                log_prob_val = None
+            try:
+                value_val = float(value.item()) if hasattr(value, 'item') else float(value)
+            except Exception:
+                value_val = None
+
+            logger.info("Step %d epoch %d | action=%s log_prob=%s value=%s reward=%.4f", step_idx, epoch, int(action.item()), log_prob_val, value_val, float(reward))
+            csv_writer.writerow([
+                datetime.utcnow().isoformat(), epoch, step_idx,
+                _truncate(query, 800), int(action.item()), log_prob_val, value_val, float(reward),
+                _truncate(context, 1000), _truncate(answer, 1000), '', ''
+            ])
+            
+            # Store experience
+            states.append(state_emb)
+            actions.append(action)
+            log_probs.append(log_prob)
+            rewards.append(reward)
+            values.append(value)
+            
+        # 2. Compute Returns and Advantages
+        # (Simplified for single-step routing: Return = Reward)
+        returns = torch.tensor(rewards, dtype=torch.float32).to(states[0].device)
+        values_tensor = torch.cat(values).squeeze()
+        advantages = returns - values_tensor.detach()
+        
+        # Normalize advantages for training stability
+        if advantages.numel() > 1:
+            advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+        else:
+            advantages = advantages - advantages.mean()
+        
+        # 3. PPO Update Phase
+        # states_tensor = torch.stack(states)
+        # actions_tensor = torch.stack(actions)
+        # log_probs_tensor = torch.stack(log_probs)
+        states_tensor = torch.stack(states).detach()
+        actions_tensor = torch.stack(actions).detach()
+        log_probs_tensor = torch.stack(log_probs).detach() 
+        
+        # Ensure these are also detached (just to be safe)
+        returns = returns.detach()
+        advantages = advantages.detach()
+        
+        
+        # loss = ppo_update(
+        #     model, optimizer, states_tensor, actions_tensor, 
+        #     log_probs_tensor, returns, advantages
+        # )
+        
+        # NEW: Run multiple optimization passes over the same rollout data
+        ppo_inner_epochs = 4
+        for _ in range(ppo_inner_epochs):
+            loss = ppo_update(
+                model, optimizer, states_tensor, actions_tensor, 
+                log_probs_tensor, returns, advantages
+            )
+        
+
+        avg_reward = float(returns.mean())
+        logger.info("PPO Epoch %d | Avg Reward: %.4f | Policy Loss: %.4f", epoch+1, avg_reward, loss)
+        csv_writer.writerow([
+            datetime.utcnow().isoformat(), epoch, '', '', '', '', '', '', '', '', avg_reward, loss
+        ])
+    
+    try:
+        csv_file.flush()
+        csv_file.close()
+    except Exception:
+        pass
+        
+    print(f"PPO Epoch {epoch+1} | Avg Reward: {returns.mean():.4f} | Policy Loss: {loss:.4f}")
+    
+    
+def run_ppo_rollout(model, embedder, queries, optimizer, epochs=40, device: str = "cpu"):
+    model.train()
+    
+    # Prepare CSV debug log (repo-level `logs/ppo_debug.csv`)
+    repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
+    csv_path = os.path.abspath(os.path.join(repo_root, 'logs', 'ppo_debug.csv'))
+    os.makedirs(os.path.dirname(csv_path), exist_ok=True)
+    write_header = not os.path.exists(csv_path)
+    csv_file = open(csv_path, mode='a', newline='', encoding='utf-8')
+    csv_writer = csv.writer(csv_file)
+    if write_header:
+        csv_writer.writerow([
+            'timestamp', 'epoch', 'step', 'query', 'action', 'action_log_prob', 'value', 'reward', 'context', 'answer', 'epoch_avg_reward', 'policy_loss'
+        ])
+    
+    for epoch in range(epochs):
+        states, actions, log_probs, rewards, values = [], [], [], [], []
+        
+        # 1. Rollout (Data Gathering)
+        # for query in queries:
+        for step_idx, query in enumerate(queries):
+            with torch.no_grad():
+                state_emb = embedder.encode(query, convert_to_tensor=True)
+            
+            # Clone and move to device
+            try:
+                state_emb = state_emb.clone().to(device)
+            except Exception:
+                state_emb = state_emb.clone().to(next(model.parameters()).device)
+                
+            action, log_prob, value = model.get_action(state_emb)
+            context, answer = simulate_rag_pipeline(query, action.item())
+            # reward = llm_judge_reward(query, context, answer, action.item())
+            reward = llm_judge_reward_openai(query, context, answer, action.item())
+            
+            # Log per-step details to logger and CSV for debugging
+            try:
+                log_prob_val = float(log_prob.item()) if hasattr(log_prob, 'item') else float(log_prob)
+            except Exception:
+                log_prob_val = None
+            try:
+                value_val = float(value.item()) if hasattr(value, 'item') else float(value)
+            except Exception:
+                value_val = None
+                
+            logger.info("Step %d epoch %d | action=%s log_prob=%s value=%s reward=%.4f", step_idx, epoch, int(action.item()), log_prob_val, value_val, float(reward))
+            csv_writer.writerow([
+                datetime.utcnow().isoformat(), epoch, step_idx,
+                _truncate(query, 800), int(action.item()), log_prob_val, value_val, float(reward),
+                _truncate(context, 1000), _truncate(answer, 1000), '', ''
+            ])
+            
+            states.append(state_emb)
+            actions.append(action)
+            log_probs.append(log_prob)
+            rewards.append(reward)
+            values.append(value)
+            
+            
+        returns = torch.tensor(rewards, dtype=torch.float32).to(states[0].device)
+        values_tensor = torch.cat(values).squeeze()
+        advantages = returns - values_tensor.detach()
+        
+        # ---------------------------------------------------------
+        # CRITICAL FIX: Strict Advantage Normalization
+        # Prevents positive-only rewards from causing gradient collapse
+        # ---------------------------------------------------------
+        # advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
+        if advantages.numel() > 1:
+            advantages = (advantages - advantages.mean()) / (advantages.std(unbiased=False) + 1e-8)
+        else:
+            advantages = advantages - advantages.mean()
+        
+        states_tensor = torch.stack(states).detach()
+        actions_tensor = torch.stack(actions).detach()
+        log_probs_tensor = torch.stack(log_probs).detach()
+        returns = returns.detach()
+        advantages = advantages.detach()
+        
+        # 2. PPO Inner Optimization Loop
+        ppo_inner_epochs = 4 
+        epoch_loss = 0
+        
+        for _ in range(ppo_inner_epochs):
+            loss, current_ent_coef = ppo_update(
+                model, optimizer, states_tensor, actions_tensor, 
+                log_probs_tensor, returns, advantages,
+                current_epoch=epoch, total_epochs=epochs
+            )
+            epoch_loss += loss
+        avg_reward = float(returns.mean())
+        logger.info("PPO Epoch %d | Avg Reward: %.4f | Policy Loss: %.4f", epoch+1, avg_reward, loss)
+        csv_writer.writerow([
+            datetime.utcnow().isoformat(), epoch, '', '', '', '', '', '', '', '', avg_reward, loss
+        ])
+            
+        print(f"Epoch {epoch+1}/{epochs} | Avg Reward: {returns.mean():.4f} | Ent Coef: {current_ent_coef:.4f} | Loss: {epoch_loss/ppo_inner_epochs:.4f}")

router_checkpoint.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f751cb17c9b905d0476abe192990de3e154d8684d5321eca08338829d72ce14c
+size 270451

router_checkpoint_after_grpo.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35126ed03ef000e4bba3b8c938b4dc9c9e375612e21ba0c60374e3ed164fd4b4
+size 270777

router_checkpoint_pt_after_ppo.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe76dc575e8b36fe20b19309c060d0f9df5518dc5afe537b2cb8c3e267629548
+size 270813

router_checkpoint_pt_after_ppo_40_epoch_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4d2fd917560dfe4398de00335dca1033614049029aa6be6ea55f8fc3be580c9
+size 270813

run_grpo.py

@@ -0,0 +1,201 @@
+#!/usr/bin/env python3
+"""
+run_grpo.py — GRPO training entry point for the RAG routing agent.
+
+This is the GRPO equivalent of run_ppo-2.py.
+
+Key differences from PPO:
+  - Uses run_grpo_rollout() instead of run_ppo_rollout()
+  - group_size (G) controls how many action samples are drawn per query
+    to compute within-group advantage normalisation (no value network needed)
+  - kl_coef controls the KL divergence penalty against the reference policy
+  - Checkpoint saved to router_checkpoint_after_grpo.pt
+"""
+import os
+# from shutil import copy
+import sys
+import copy
+# from asyncssh import logger
+import torch
+import pandas as pd
+
+# Ensure the routing_agent package is importable
+ROOT = os.path.dirname(__file__)
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from grpo_trainer import run_grpo_rollout, load_router_policy
+from model import RouterNetwork
+from sentence_transformers import SentenceTransformer
+
+
+def main_from_scratch():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # ------------------------------------------------------------------ #
+    # Embedder + Model
+    # ------------------------------------------------------------------ #
+    embed_model_name = "all-MiniLM-L6-v2"
+    embedder = SentenceTransformer(embed_model_name)
+    embed_dim = embedder.get_sentence_embedding_dimension()
+
+    model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)
+
+    # ------------------------------------------------------------------ #
+    # Optionally resume from an existing checkpoint
+    # ------------------------------------------------------------------ #
+    checkpoint_path = os.path.join(os.path.dirname(__file__), "router_checkpoint.pt")
+    if os.path.exists(checkpoint_path):
+        try:
+            sd = torch.load(checkpoint_path, map_location="cpu")
+            if isinstance(sd, dict) and "model_state_dict" in sd:
+                model.load_state_dict(sd["model_state_dict"], strict=False)
+            else:
+                model.load_state_dict(sd, strict=False)
+            print(f"Loaded checkpoint: {checkpoint_path}")
+        except Exception as e:
+            print(f"Failed loading checkpoint: {e}")
+    else:
+        print("No existing checkpoint found — training from scratch.")
+
+    # ------------------------------------------------------------------ #
+    # Optimizer
+    # ------------------------------------------------------------------ #
+    optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)
+
+    # ------------------------------------------------------------------ #
+    # Training queries
+    # ------------------------------------------------------------------ #
+    df = pd.read_csv("df_queries.csv")
+    training_queries = df["question"].tolist()
+
+    # Uncomment to limit queries for a quick smoke-test:
+    # training_queries = training_queries[:2]
+
+    print(f"Training on {len(training_queries)} queries.")
+
+    # ------------------------------------------------------------------ #
+    # GRPO Hyperparameters
+    # ------------------------------------------------------------------ #
+    EPOCHS     = 20#40       # Total training epochs (same as PPO run)
+    GROUP_SIZE = 4        # G: number of action samples drawn per query per epoch
+                          # Higher G → better advantage estimates, but more API calls.
+                          # Set to 1 to match the original PPO single-sample behaviour.
+    KL_COEF    = 0.04     # Weight for KL divergence penalty (DeepSeek-R1 default: 0.04)
+    CLIP_RATIO = 0.2      # PPO-style clipping kept in GRPO surrogate
+
+    # ------------------------------------------------------------------ #
+    # Run GRPO rollouts
+    # ------------------------------------------------------------------ #
+    # run_grpo_rollout_old(
+    #     model=model,
+    #     embedder=embedder,
+    #     queries=training_queries,
+    #     optimizer=optimizer,
+    #     epochs=EPOCHS,
+    #     group_size=GROUP_SIZE,
+    #     device=str(device),
+    #     kl_coef=KL_COEF,
+    #     clip_ratio=CLIP_RATIO,
+    # )
+    run_grpo_rollout(
+        model=model,
+        embedder=embedder,
+        queries=training_queries,
+        optimizer=optimizer,
+        epochs=EPOCHS,
+        device=str(device),
+        kl_coef=KL_COEF,
+        clip_ratio=CLIP_RATIO,
+    )
+
+    # ------------------------------------------------------------------ #
+    # Save updated checkpoint
+    # ------------------------------------------------------------------ #
+    out_path = os.path.join(os.path.dirname(__file__), "router_checkpoint_after_grpo.pt")
+    torch.save({"model_state_dict": model.state_dict()}, out_path)
+    print(f"Saved GRPO-updated checkpoint to: {out_path}")
+
+def main():
+    LR = 3e-5
+    # ------------------------------------------------------------------ #
+    # GRPO Hyperparameters
+    # ------------------------------------------------------------------ #
+    EPOCHS     = 20#40       # Total training epochs (same as PPO run)
+    GROUP_SIZE = 4        # G: number of action samples drawn per query per epoch
+                          # Higher G → better advantage estimates, but more API calls.
+                          # Set to 1 to match the original PPO single-sample behaviour.
+    KL_COEF    = 0.01#0.04     # Weight for KL divergence penalty (DeepSeek-R1 default: 0.04)
+    #Your first run showed KL was essentially zero for the first 6 epochs — meaning kl_coef=0.04 was over-constraining the policy and slowing early learning. With the reference refresh now in place (point 1 above), you can lower this safely:
+    
+    CLIP_RATIO = 0.2      # PPO-style clipping kept in GRPO surrogate
+    
+    RESUME_FROM = "adapters/router_grpo_epoch_14.pt"  # set to None to start fresh
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # EMBEDDER (always needed to encode queries)
+    embed_model_name = "all-MiniLM-L6-v2"
+    embedder = SentenceTransformer(embed_model_name)
+    embed_dim = embedder.get_sentence_embedding_dimension()
+    
+    # ------------------------------------------------------------------ #
+    # MODEL — build the architecture shell only (no SFT loading)
+    # ------------------------------------------------------------------ #
+    model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+    
+
+    start_epoch = 0
+    if RESUME_FROM and os.path.exists(RESUME_FROM):
+        ckpt = torch.load(RESUME_FROM, map_location=device)
+        model.load_state_dict(ckpt["model_state_dict"])
+        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+        start_epoch = ckpt["epoch"]          # = 7 → loop starts at epoch 7
+        print(f"Resumed from {RESUME_FROM} — continuing from epoch {start_epoch}")
+    else:
+        print(f"No checkpoint found at {RESUME_FROM}, starting from scratch.")
+    
+    # ------------------------------------------------------------------ #
+    # REF POLICY — snapshot AFTER loading checkpoint, BEFORE training loop
+    # This is the fixed reference the KL penalty measures distance from.
+    # ------------------------------------------------------------------ #
+    ref_policy = copy.deepcopy(model)
+    ref_policy.eval()
+    for p in ref_policy.parameters():
+        p.requires_grad_(False)    # freeze completely — never updated
+        
+    # ------------------------------------------------------------------ #
+    # Training queries
+    # ------------------------------------------------------------------ #
+    df = pd.read_csv("df_queries_filtered.csv")
+    training_queries = df["question"].tolist()
+
+    # Uncomment to limit queries for a quick smoke-test:
+    # training_queries = training_queries[:2]
+
+    print(f"Training on {len(training_queries)} queries.")
+
+    run_grpo_rollout(
+        model=model,
+        embedder=embedder,
+        queries=training_queries,
+        optimizer=optimizer,
+        ref_policy   = ref_policy,      # ← passed in, not re-created inside trainer
+        epochs=EPOCHS,
+        start_epoch=start_epoch,
+        device=str(device),
+        kl_coef=KL_COEF,
+        clip_ratio=CLIP_RATIO,
+    )
+
+    # ------------------------------------------------------------------ #
+    # Save updated checkpoint
+    # ------------------------------------------------------------------ #
+    out_path = os.path.join(os.path.dirname(__file__), "router_checkpoint_after_grpo.pt")
+    torch.save({"model_state_dict": model.state_dict()}, out_path)
+    print(f"Saved GRPO-updated checkpoint to: {out_path}")
+
+if __name__ == "__main__":
+    main()

run_inference.py

@@ -0,0 +1,112 @@
+import os
+import sys
+import time
+import torch
+from sentence_transformers import SentenceTransformer
+import pandas as pd
+
+# ── 1. Path setup (same pattern as run_ppo-2.py) ────────────────────────────
+ROOT = os.path.dirname(__file__)
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from model import RouterNetwork
+from grpo_trainer import simulate_rag_pipeline   # reuse your existing tool caller
+
+# ── 2. Config ────────────────────────────────────────────────────────────────
+CHECKPOINT  = os.path.join(ROOT, "adapters/router_grpo_epoch_16.pt")   # ← change epoch here
+EMBED_MODEL = "all-MiniLM-L6-v2"
+DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"
+
+ACTION_NAMES = {
+    0: "Guidelines  (academic rules / visa / scholarships)",
+    1: "Web-Search   (news / announcements)",
+    2: "CBNU Database (courses / departments / contacts)",
+}
+
+# ── 3. Load embedder + model ─────────────────────────────────────────────────
+embedder  = SentenceTransformer(EMBED_MODEL)
+embed_dim = embedder.get_sentence_embedding_dimension()
+model     = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(DEVICE)
+
+# Load checkpoint (supports both raw state_dict and wrapped dict)
+sd = torch.load(CHECKPOINT, map_location="cpu")
+if isinstance(sd, dict) and "model_state_dict" in sd:
+    model.load_state_dict(sd["model_state_dict"], strict=False)
+else:
+    model.load_state_dict(sd, strict=False)
+
+model.eval()
+print(f"✅  Loaded checkpoint: {CHECKPOINT}")
+print(f"    Device: {DEVICE}\n")
+
+# ── 4. Single-query inference function ───────────────────────────────────────
+def route_query(query: str, call_tool: bool = True):
+    """
+    Route one query through the loaded GRPO checkpoint.
+
+    Args:
+        query     : The user question string.
+        call_tool : If True, also calls simulate_rag_pipeline and prints the answer.
+                    Set False for routing-only / fast testing.
+
+    Returns:
+        action_idx (int), action_name (str), context (str), answer (str)
+    """
+    # Embed
+    with torch.no_grad():
+        emb = embedder.encode(query, convert_to_tensor=True).to(DEVICE)
+        action, log_prob, _ = model.get_action(emb)
+        action_idx = int(action.item())
+
+    action_name = ACTION_NAMES[action_idx]
+    print(f"Query   : {query}")
+    print(f"→ Routed to Action {action_idx}: {action_name}")
+    print(f"  log_prob = {log_prob.item():.4f}")
+
+    context, answer = "", ""
+    if call_tool:
+        print("  Calling tool …")
+        context, answer = simulate_rag_pipeline(query, action_idx)
+        print(f"  Answer  : {answer[:300]}")
+
+    return action_idx, action_name, context, answer
+
+
+# ── 5. Main: test with one or more queries ───────────────────────────────────
+if __name__ == "__main__":
+    test_queries = [
+        "What are the graduation requirements for international students?",
+        "How do I register for courses next semester?",
+        "What is the latest news from CBNU?",
+    ]
+    test_queries = [
+        "2학기 개강일은?",
+        "알바하고 싶은데 어떤 조건이 필요한가?",
+        "아르바이트 하려는데 조건은?",
+        "시간제 취업 방법은?",
+        "건강보험 어떻게 가입해?",
+        "성적확인 기간은?"
+    ]
+    
+    eval_df = pd.read_csv('ko_eval_queries.csv')
+    test_queries = eval_df['question'].tolist()
+    results = []
+    for idx, row in eval_df.iterrows():
+        q = row.get('question')
+        ground = row.get('action_label')
+        start = time.perf_counter()
+        pred_idx, pred_name, context, answer = route_query(q, call_tool=False)
+        end = time.perf_counter()
+        elapsed = end - start
+        results.append({
+            'question': q,
+            'ground_truth': ground,
+            'predicted_label': pred_idx,
+            'time_took': elapsed,
+        })
+        
+        
+    out_df = pd.DataFrame(results)
+    out_df.to_csv('ko_eval_with_timing_results_epoch_16.csv', index=False)
+    print(f"Wrote results to eval_with_timing_results.csv ({len(out_df)} rows)")

run_ppo.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+import os
+import sys
+import torch
+import pandas as pd
+
+# Ensure routing_agent package is importable
+ROOT = os.path.dirname(__file__)
+if ROOT not in sys.path:
+    sys.path.insert(0, ROOT)
+
+from ppo_trainer import run_ppo_rollout, load_router_policy
+from model import RouterNetwork
+from sentence_transformers import SentenceTransformer
+
+
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Prepare embedder and model
+    embed_model_name = 'all-MiniLM-L6-v2'
+    embedder = SentenceTransformer(embed_model_name)
+
+    embed_dim = embedder.get_sentence_embedding_dimension()
+    model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)
+
+    # Optionally load existing checkpoint to continue training
+    checkpoint_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint.pt')
+    if os.path.exists(checkpoint_path):
+        try:
+            sd = torch.load(checkpoint_path, map_location='cpu')
+            if isinstance(sd, dict) and 'model_state_dict' in sd:
+                model.load_state_dict(sd['model_state_dict'], strict=False)
+            else:
+                model.load_state_dict(sd, strict=False)
+            print('Loaded checkpoint', checkpoint_path)
+        except Exception as e:
+            print('Failed loading checkpoint:', e)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)
+
+    # training_queries = [
+    #     "What are the graduation requirements for the 2026 academic year?",
+    #     "Show me the syllabus for the Advanced Database course.",
+    #     "How do I apply for the on-campus dormitory?"
+    # ]
+    df = pd.read_csv('df_queries.csv')
+    training_queries = df['question'].tolist()
+    
+    # training_queries = training_queries[:2]  # Limit to first 200 queries for faster testing
+
+
+    # Run PPO rollouts (this uses simulate_rag_pipeline to call the agent tools)
+    run_ppo_rollout(model, embedder, training_queries, optimizer, epochs=40, device=str(device))
+
+    # Save the updated model
+    out_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint_pt_after_ppo.pt')
+    torch.save({'model_state_dict': model.state_dict()}, out_path)
+    print('Saved PPO-updated checkpoint to', out_path)
+
+
+if __name__ == '__main__':
+    main()

sft_trainer.py

@@ -0,0 +1,46 @@
+import torch
+import torch.optim as optim
+import torch.nn as nn
+from model import RouterNetwork
+
+def train_sft(model, embedder, dataset, epochs = 5, lr = 1e-3):
+    """
+    dataset: List of dicts [{'question': '...', 'label': 0|1|2}]
+    """
+    device = next(model.parameters()).device
+    class_weights = torch.tensor([1.0, 1.0, 10.0], device=device, dtype=torch.float)
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    criterion = nn.CrossEntropyLoss(weight=class_weights)
+    
+    model.train()
+    for epoch in range(epochs):
+        total_loss = 0
+        correct_predictions = 0
+        for data in dataset:
+            # 1. Embed the query (state)
+            # Using torch.no_grad() for embedder so we only train the router MLP
+            with torch.no_grad():
+                query_emb = embedder.encode(data['question'], convert_to_tensor=True)
+                
+            label = torch.tensor([data['label']], dtype=torch.long).to(query_emb.device)
+            
+            # 2. Forward pass through the router network
+            optimizer.zero_grad()
+            action_logits, _ = model(query_emb) # only care about action_logits (Actor head) for SFT
+
+            # Calculate prediction and track accuracy for this epoch
+            predicted_class = torch.argmax(action_logits, dim=-1)
+            if predicted_class.item() == label.item():
+                correct_predictions += 1
+            
+            # 3. Compute loss and backpropagate
+            # CrossEntropyLoss expects (N, C) for logits and (N) for labels, so we unsqueeze action_logits to add batch dimension
+            loss = criterion(action_logits.unsqueeze(0), label)
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+            
+        avg_loss = total_loss / len(dataset)
+        accuracy = (correct_predictions / len(dataset)) * 100
+        print(f"SFT Epoch {epoch+1} | Loss: {avg_loss:.4f} | Accuracy: {accuracy:.2f}%")