Uploaded train.py

train.py

@@ -1,840 +1,860 @@
-"""
-train.py — LogTriageEnv GRPO Training Loop
-Meta × PyTorch × Scaler OpenEnv Hackathon — Grand Finale
-
-Usage:
-    python train.py --model HuggingFaceTB/SmolLM2-360M-Instruct --task single_crash --episodes 50 --env_url http://localhost:7860
-    python train.py --model HuggingFaceTB/SmolLM2-360M-Instruct --task all --episodes 100 --env_url http://localhost:7860
-
-    # Colab T4 GPU — use Unsloth (recommended for Qwen 3B/7B):
-    python train.py --model Qwen/Qwen2.5-7B-Instruct --task all --episodes 50 --use_unsloth --env_url https://ogrohit-logtriage-env.hf.space
-    python train.py --model Qwen/Qwen2.5-3B-Instruct --task all --episodes 50 --use_unsloth --env_url https://ogrohit-logtriage-env.hf.space
-
-    # Local laptop (no quantization):
-    python train.py --model HuggingFaceTB/SmolLM2-360M-Instruct --task all --episodes 50 --env_url http://localhost:7860
-
-    # Onsite with A100 — use Unsloth for max speed:
-    python train.py --model Qwen/Qwen2.5-32B-Instruct --task all --episodes 100 --use_unsloth --env_url https://ogrohit-logtriage-env.hf.space
-"""
-
-import argparse
-import json
-import re
-import time
-import os
-from dataclasses import dataclass, field
-from typing import Optional, List
-
-import requests
-import matplotlib.pyplot as plt
-import matplotlib
-matplotlib.use("Agg")  # headless — no display required
-
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-from trl import GRPOConfig, GRPOTrainer
-from datasets import Dataset
-
-try:
-    from peft import LoraConfig, get_peft_model, PeftModel
-    PEFT_AVAILABLE = True
-except ImportError:
-    PEFT_AVAILABLE = False
-
-try:
-    from unsloth import FastLanguageModel
-    UNSLOTH_AVAILABLE = True
-except ImportError:
-    UNSLOTH_AVAILABLE = False
-
-# ── Constants ────────────────────────────────────────────────────────────────
-
-VALID_ACTION_TYPES = [
-    "classify_severity",
-    "identify_root_cause",
-    "escalate",
-    "remediate",
-    "request_more_logs",
-    "resolve",
-    "ignore",
-]
-
-VALID_VALUES = {
-    "classify_severity": ["P1", "P2", "P3"],
-    "identify_root_cause": [
-        "api-gateway", "auth-service", "user-db",
-        "payment-service", "payment-db",
-        "notification-service", "email-queue",
-    ],
-    "escalate": ["sre-team", "backend-team", "dba-team", "security-team", "ignore"],
-    "remediate": [
-        "restart:api-gateway", "restart:auth-service", "restart:user-db",
-        "restart:payment-service", "restart:payment-db",
-        "restart:notification-service", "restart:email-queue",
-        "rollback:api-gateway", "rollback:auth-service", "rollback:payment-service",
-        "scale:api-gateway", "scale:payment-service",
-        "flush-cache:user-db", "flush-cache:payment-db",
-        "kill-query:user-db", "kill-query:payment-db",
-    ],
-    "request_more_logs": [
-        "api-gateway", "auth-service", "user-db",
-        "payment-service", "payment-db",
-        "notification-service", "email-queue", "all",
-    ],
-    "resolve": ["resolved"],
-    "ignore": ["noise"],
-}
-
-SYSTEM_PROMPT = """You are an expert SRE (Site Reliability Engineer) triaging a live production incident.
-
-You will receive log lines from a microservice cluster. Your job is to reason carefully and take ONE action per step.
-
-The service topology is:
-  [api-gateway] → [auth-service] → [user-db]
-               → [payment-service] → [payment-db]
-               → [notification-service] → [email-queue]
-
-Available actions:
-- classify_severity: Set priority. Values: P1 (customer-facing outage), P2 (degradation), P3 (warning)
-- identify_root_cause: Point to the failing service. Values: api-gateway, auth-service, user-db, payment-service, payment-db, notification-service, email-queue
-- escalate: Page a team. Values: sre-team, backend-team, dba-team, security-team, ignore
-- remediate: Apply a fix. Values: restart:<service>, rollback:<service>, scale:<service>, flush-cache:<service>, kill-query:<service>
-- request_more_logs: Get more logs. Values: <service-name> or all
-- resolve: Mark incident resolved. Value: resolved
-- ignore: Mark as noise. Value: noise
-
-CRITICAL RULES:
-1. For cascading failures, find the ROOT CAUSE service, not the first service that shows errors
-2. P1 = customer-facing impact (error rate >5%), P2 = degradation, P3 = warning only
-3. Do NOT over-escalate. Paging the wrong team is penalized.
-4. Be efficient — unnecessary steps reduce your score.
-
-You MUST respond in this exact JSON format and nothing else:
-{
-  "action_type": "<one of the action types above>",
-  "value": "<valid value for that action type>",
-  "confidence": <float 0.0-1.0>,
-  "reasoning": "<one sentence explaining why>"
-}"""
-
-
-# ── Env Client ──────────────────────────────────────────────���────────────────
-
-class LogTriageEnvClient:
-    """HTTP client for LogTriageEnv."""
-
-    def __init__(self, base_url: str):
-        self.base_url = base_url.rstrip("/")
-        self._verify_connection()
-
-    def _verify_connection(self):
-        try:
-            r = requests.get(f"{self.base_url}/health", timeout=10)
-            r.raise_for_status()
-            print(f"[OK] Connected to LogTriageEnv at {self.base_url}")
-        except Exception as e:
-            raise RuntimeError(
-                f"[ERROR] Cannot reach LogTriageEnv at {self.base_url}\n"
-                f"   Make sure Docker is running: docker run -p 7860:7860 logtriage-env\n"
-                f"   Error: {e}"
-            )
-
-    def reset(self, task_id: str, seed: int = 42) -> dict:
-        r = requests.post(
-            f"{self.base_url}/reset",
-            json={"task_id": task_id, "seed": seed},
-            timeout=15,
-        )
-        r.raise_for_status()
-        return r.json()
-
-    def step(self, action: dict) -> dict:
-        r = requests.post(
-            f"{self.base_url}/step",
-            json=action,
-            timeout=15,
-        )
-        r.raise_for_status()
-        return r.json()
-
-    def get_tasks(self) -> list:
-        r = requests.get(f"{self.base_url}/tasks", timeout=10)
-        r.raise_for_status()
-        return r.json()["tasks"]
-
-
-# ── Observation Formatting ───────────────────────────────────────────────────
-
-def format_observation(obs: dict, step: int) -> str:
-    """Convert raw env observation dict into a clean prompt string."""
-    lines = []
-
-    lines.append(f"=== INCIDENT TRIAGE — Step {step} ===")
-    lines.append(f"Incident ID: {obs.get('incident_id', 'unknown')}")
-    lines.append(f"Active Alerts: {', '.join(obs.get('active_alerts', []))}")
-    lines.append("")
-
-    # System state
-    lines.append("--- System State ---")
-    system_state = obs.get("system_state", {})
-    for svc, status in system_state.items():
-        if isinstance(status, dict):
-            lines.append(
-                f"  {svc}: {status.get('status','?')} | "
-                f"error_rate={status.get('error_rate', 0):.1%} | "
-                f"p99={status.get('latency_p99_ms', 0)}ms"
-            )
-        else:
-            lines.append(f"  {svc}: {status}")
-
-    # Log lines
-    lines.append("")
-    lines.append("--- Log Stream ---")
-    logs = obs.get("logs", [])
-    if isinstance(logs, list):
-        for log in logs[-15:]:  # last 15 lines to stay within context
-            if isinstance(log, dict):
-                ts = log.get("timestamp", "")
-                level = log.get("level", "")
-                svc = log.get("service", "")
-                msg = log.get("message", "")
-                lines.append(f"  [{ts}] {level:5} {svc:25} {msg}")
-            else:
-                lines.append(f"  {log}")
-    else:
-        lines.append(str(logs))
-
-    # Feedback from last action
-    feedback = obs.get("last_action_feedback", "")
-    if feedback:
-        lines.append("")
-        lines.append(f"--- Last Action Feedback ---")
-        lines.append(f"  {feedback}")
-
-    lines.append("")
-    lines.append("What is your next action? Respond in JSON only.")
-
-    return "\n".join(lines)
-
-
-# ── Action Parsing ────────────────────────────────────────────────────────────
-
-def parse_action(llm_output: str) -> Optional[dict]:
-    """
-    Parse LLM output into a valid TriageAction dict.
-    Returns None if parsing fails completely.
-    """
-    # Try direct JSON parse first
-    try:
-        # Strip markdown code fences if present
-        clean = re.sub(r"```(?:json)?", "", llm_output).strip().rstrip("```").strip()
-        # Find first { ... } block
-        match = re.search(r"\{.*\}", clean, re.DOTALL)
-        if match:
-            action = json.loads(match.group())
-            if "action_type" in action and "value" in action:
-                # Validate action_type
-                if action["action_type"] not in VALID_ACTION_TYPES:
-                    return None
-                # Validate value against strict server-side rules
-                validated = _validate_action_value(action["action_type"], action.get("value", ""))
-                if validated is None:
-                    return None
-                action["value"] = validated
-                action["confidence"] = 0.5
-                action["reasoning"] = ""
-                return action
-    except (json.JSONDecodeError, KeyError):
-        pass
-
-    # Fallback: keyword extraction (only on known-good pairs)
-    output_lower = llm_output.lower()
-    for action_type in VALID_ACTION_TYPES:
-        if action_type.replace("_", " ") in output_lower or action_type in output_lower:
-            for value in VALID_VALUES.get(action_type, []):
-                if value.lower() in output_lower:
-                    # Extra validation for escalate: "ignore" is NOT a valid escalate value
-                    if action_type == "escalate" and value == "ignore":
-                        continue
-                    return {
-                        "action_type": action_type,
-                        "value": value,
-                        "confidence": 0.3,
-                        "reasoning": "parsed via fallback",
-                    }
-
-    # Last resort: safe default
-    return {
-        "action_type": "request_more_logs",
-        "value": "all",
-        "confidence": 0.1,
-        "reasoning": "failed to parse LLM output",
-    }
-
-
-def _validate_action_value(action_type: str, value: str) -> Optional[str]:
-    """Validate action value against server-side rules. Returns clean value or None."""
-    if action_type == "classify_severity":
-        if value in ("P1", "P2", "P3"):
-            return value
-    elif action_type == "identify_root_cause":
-        valid = {
-            "api-gateway", "auth-service", "user-db",
-            "payment-service", "payment-db",
-            "notification-service", "email-queue",
-        }
-        if value in valid:
-            return value
-        # Fuzzy match: "payment" -> "payment-service"
-        if value in ("payment", "payment svc", "paymentservice"):
-            return "payment-service"
-        if value in ("user", "userdb", "user_db"):
-            return "user-db"
-        if value in ("auth", "authsvc"):
-            return "auth-service"
-        if value in ("api", "gateway", "api-gw"):
-            return "api-gateway"
-        if value in ("notif", "notification", "notif-service"):
-            return "notification-service"
-        if value in ("email", "emailqueue", "queue"):
-            return "email-queue"
-    elif action_type == "escalate":
-        valid = {"sre-team", "backend-team", "dba-team", "security-team"}
-        if value in valid:
-            return value
-    elif action_type == "remediate":
-        if ":" in value:
-            prefix, service = value.split(":", 1)
-            valid_prefixes = {"restart", "rollback", "scale", "flush-cache", "kill-query"}
-            if prefix in valid_prefixes:
-                # Map service aliases
-                service_map = {
-                    "payment": "payment-service",
-                    "userdb": "user-db",
-                    "user_db": "user-db",
-                    "auth": "auth-service",
-                    "api": "api-gateway",
-                    "gateway": "api-gateway",
-                    "notif": "notification-service",
-                    "email": "email-queue",
-                }
-                clean_service = service_map.get(service, service)
-                return f"{prefix}:{clean_service}"
-    elif action_type == "request_more_logs":
-        valid_services = {
-            "api-gateway", "auth-service", "user-db",
-            "payment-service", "payment-db",
-            "notification-service", "email-queue", "all",
-        }
-        if value in valid_services:
-            return value
-        service_map = {
-            "payment": "payment-service", "userdb": "user-db",
-            "user_db": "user-db", "auth": "auth-service",
-            "api": "api-gateway", "gateway": "api-gateway",
-            "notif": "notification-service", "email": "email-queue",
-        }
-        if value in service_map:
-            return service_map[value]
-    elif action_type == "resolve":
-        if value == "resolved":
-            return "resolved"
-    elif action_type == "ignore":
-        if value == "noise":
-            return "noise"
-    return None
-
-
-# ── Single Episode Rollout ───────────────────────────────────────────────────
-
-def run_episode(
-    env: LogTriageEnvClient,
-    model,
-    tokenizer,
-    task_id: str,
-    seed: int,
-    device: str,
-    max_steps: int = 15,
-    verbose: bool = False,
-) -> tuple[float, int, list[dict]]:
-    """
-    Run one full episode.
-    Returns: (total_reward, steps_taken, trajectory)
-    trajectory = list of {prompt, response, reward} dicts for GRPO
-    """
-    obs = env.reset(task_id=task_id, seed=seed)
-    total_reward = 0.0
-    steps = 0
-    trajectory = []
-    done = False
-
-    while not done and steps < max_steps:
-        # Format observation into prompt
-        prompt_text = format_observation(obs, steps + 1)
-
-        # Build chat messages
-        messages = [
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user", "content": prompt_text},
-        ]
-
-        # Tokenize
-        input_ids = tokenizer.apply_chat_template(
-            messages,
-            return_tensors="pt",
-            add_generation_prompt=True,
-        )
-        input_ids = input_ids["input_ids"].to(device)
-        attention_mask = (input_ids != tokenizer.pad_token_id).long()
-        gen_kwargs = {
-            "max_new_tokens": 150,
-            "do_sample": True,
-            "temperature": 0.7,
-            "top_p": 0.9,
-            "attention_mask": attention_mask,
-            "pad_token_id": tokenizer.eos_token_id,
-        }
-
-        # Generate
-        with torch.no_grad():
-            output_ids = model.generate(input_ids, **gen_kwargs)
-
-        # Decode only the new tokens
-        prompt_len = input_ids.shape[1]
-        new_tokens = output_ids[0][prompt_len:]
-        llm_output = tokenizer.decode(new_tokens, skip_special_tokens=True)
-
-        # Parse action
-        action = parse_action(llm_output)
-        if action is None:
-            action = {"action_type": "request_more_logs", "value": "all",
-                      "confidence": 0.1, "reasoning": "parse failed"}
-
-        # Step env
-        try:
-            obs = env.step(action)
-        except requests.HTTPError as e:
-            if verbose:
-                print(f"[WARN] Step HTTP error: {e}")
-            break
-
-        # Extract reward
-        step_reward = obs.get("reward", 0.0)
-        total_reward += step_reward
-        done = obs.get("done", False)
-        steps += 1
-
-        # Store for GRPO
-        trajectory.append({
-            "prompt": prompt_text,
-            "response": llm_output,
-            "reward": step_reward,
-        })
-
-        if verbose:
-            print(f"    Step {steps}: action={action['action_type']}({action['value']}) "
-                  f"reward={step_reward:+.2f} done={done}")
-
-    return total_reward, steps, trajectory
-
-
-# ── Reward Curve Plot ─────────────────────────────────────────────────────────
-
-def save_reward_curve(history: dict[str, list[float]], output_path: str = "reward_curve.png"):
-    """
-    history: {"single_crash": [r1, r2, ...], "cascading_failure": [...], ...}
-    """
-    fig, ax = plt.subplots(figsize=(10, 6))
-
-    colors = {"single_crash": "#00C49F", "cascading_failure": "#FFBB28", "silent_degradation": "#FF6B6B"}
-    labels = {"single_crash": "Task 1: Single Crash (Easy)",
-              "cascading_failure": "Task 2: Cascading Failure (Medium)",
-              "silent_degradation": "Task 3: Silent Degradation (Hard)"}
-
-    for task_id, rewards in history.items():
-        if not rewards:
-            continue
-        # Smooth with rolling average (window=5)
-        smoothed = []
-        for i in range(len(rewards)):
-            window = rewards[max(0, i-4):i+1]
-            smoothed.append(sum(window) / len(window))
-
-        episodes = list(range(1, len(rewards) + 1))
-        color = colors.get(task_id, "#8884d8")
-        label = labels.get(task_id, task_id)
-
-        ax.plot(episodes, rewards, alpha=0.3, color=color, linewidth=0.8)
-        ax.plot(episodes, smoothed, color=color, linewidth=2.5, label=label)
-
-    ax.set_xlabel("Episode", fontsize=13)
-    ax.set_ylabel("Episode Reward", fontsize=13)
-    ax.set_title("LogTriageEnv — Agent Reward Improvement During GRPO Training", fontsize=14, fontweight="bold")
-    ax.legend(fontsize=11)
-    ax.grid(True, alpha=0.3)
-    ax.set_ylim(bottom=0)
-
-    # Add annotation
-    ax.annotate(
-        "Higher = agent solves incidents faster with fewer wrong actions",
-        xy=(0.02, 0.02), xycoords="axes fraction",
-        fontsize=9, color="gray", style="italic"
-    )
-
-    plt.tight_layout()
-    plt.savefig(output_path, dpi=150, bbox_inches="tight")
-    plt.close()
-    print(f"[PLOT] Reward curve saved -> {output_path}")
-
-
-# ── GRPO Dataset Builder ──────────────────────────────────────────────────────
-
-def build_grpo_dataset(trajectories: list[dict]) -> Dataset:
-    """
-    Build a HF Dataset from collected trajectories for GRPOTrainer.
-    Format: {"prompt": str, "completion": str, "reward": float}
-    """
-    if not trajectories:
-        # Return minimal dummy dataset if no trajectories yet
-        return Dataset.from_dict({
-            "prompt": ["dummy"],
-            "completion": ["{}"],
-            "reward": [0.0],
-        })
-
-    return Dataset.from_dict({
-        "prompt": [t["prompt"] for t in trajectories],
-        "completion": [t["response"] for t in trajectories],
-        "reward": [t["reward"] for t in trajectories],
-    })
-
-
-# ── Main Training Loop ────────────────────────────────────────────────────────
-
-def main():
-    parser = argparse.ArgumentParser(description="LogTriageEnv GRPO Training")
-    parser.add_argument("--model", default="HuggingFaceTB/SmolLM2-360M-Instruct",
-                        help="HuggingFace model ID")
-    parser.add_argument("--task", default="single_crash",
-                        choices=["single_crash", "cascading_failure", "silent_degradation", "all"],
-                        help="Task to train on. 'all' trains on all 3.")
-    parser.add_argument("--episodes", type=int, default=50,
-                        help="Number of training episodes per task")
-    parser.add_argument("--env_url", default="http://localhost:7860",
-                        help="LogTriageEnv base URL")
-    parser.add_argument("--output_dir", default="./logtriage-trained",
-                        help="Where to save the trained model")
-    parser.add_argument("--push_to_hub", action="store_true",
-                        help="Push trained model to HuggingFace Hub")
-    parser.add_argument("--hub_model_id", default=None,
-                        help="HF Hub model ID (e.g. username/logtriage-sre-agent)")
-    parser.add_argument("--verbose", action="store_true",
-                        help="Print step-by-step actions during episodes")
-    parser.add_argument("--load_in_4bit", action="store_true",
-                        help="Load model with 4-bit QLoRA quantization via BitsAndBytes (for large models on limited VRAM)")
-    parser.add_argument("--use_unsloth", action="store_true",
-                        help="Load model using Unsloth (recommended for Qwen on T4/A100 — faster and more memory efficient)")
-    parser.add_argument("--skip_grpo", action="store_true",
-                        help="Skip GRPO fine-tuning and only run rollout episodes (useful when debugging or avoiding OOM)")
-    parser.add_argument("--grpo_max_steps", type=int, default=35,
-                        help="Maximum GRPO optimization steps after rollout (default: 35)")
-    args = parser.parse_args()
-
-    # ── Setup ────────────────────────────────────────────────────────────────
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print("\n[LOGGING] LogTriageEnv GRPO Training")
-    print(f"   Model:   {args.model}")
-    print(f"   Task:    {args.task}")
-    print(f"   Episodes: {args.episodes}")
-    print(f"   Device:  {device}")
-    print(f"   Env URL: {args.env_url}\n")
-
-    # Connect to env
-    env = LogTriageEnvClient(args.env_url)
-
-    # Determine tasks to train on
-    if args.task == "all":
-        tasks = ["single_crash", "cascading_failure", "silent_degradation"]
-    else:
-        tasks = [args.task]
-
-    # Load model + tokenizer
-    print(f"[MODEL] Loading model: {args.model}")
-    use_unsloth = getattr(args, "use_unsloth", False)
-    use_lora = False
-
-    # ── Unsloth Path (recommended for Qwen on T4/A100) ───────────────────────
-    if use_unsloth and device == "cuda" and UNSLOTH_AVAILABLE:
-        print("[UNSLOTH] Loading model with Unsloth...")
-        model, tokenizer = FastLanguageModel.from_pretrained(
-            model_name=args.model,
-            max_seq_length=2048,
-            load_in_4bit=True,
-            dtype=None,  # Auto-detect
-        )
-        print(f"[OK] Model loaded via Unsloth (4-bit)")
-
-        # Apply LoRA via Unsloth
-        print("[UNSLOTH] Applying LoRA adapter (r=16, alpha=32)...")
-        model = FastLanguageModel.get_peft_model(
-            model,
-            r=16,
-            lora_alpha=32,
-            target_modules=[
-                "q_proj", "k_proj", "v_proj", "o_proj",
-                "gate_proj", "up_proj", "down_proj",
-            ],
-            lora_dropout=0.05,
-            bias="none",
-        )
-        model.print_trainable_parameters()
-        use_lora = True
-        print(f"[OK] Unsloth LoRA attached")
-        print(f"[OK] Model loaded\n")
-
-    # ── BitsAndBytes QLoRA Path (manual, or fallback) ─────────────────────────
-    elif getattr(args, "load_in_4bit", False) and device == "cuda":
-        print("[QLoRA] Loading model with BitsAndBytes 4-bit...")
-        tokenizer = AutoTokenizer.from_pretrained(args.model)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.float16,
-            bnb_4bit_use_double_quant=True,
-        )
-        print(f"[OK] 4-bit BitsAndBytesConfig applied")
-
-        model = AutoModelForCausalLM.from_pretrained(
-            args.model,
-            quantization_config=bnb_config,
-            device_map="auto",
-        )
-        print(f"[OK] Model loaded in 4-bit quantized mode")
-
-        if PEFT_AVAILABLE:
-            print("[QLoRA] Applying LoRA adapter...")
-            lora_config = LoraConfig(
-                r=16,
-                lora_alpha=32,
-                target_modules=[
-                    "q_proj", "k_proj", "v_proj", "o_proj",
-                    "gate_proj", "up_proj", "down_proj",
-                ],
-                lora_dropout=0.05,
-                bias="none",
-                task_type="CAUSAL_LM",
-            )
-            model = get_peft_model(model, lora_config)
-            model.print_trainable_parameters()
-            use_lora = True
-            print(f"[OK] LoRA adapter attached (r=16, alpha=32)")
-        else:
-            print("[WARN] PEFT not installed. Using quantized model without LoRA.")
-
-        if not hasattr(model, "processing_class"):
-            model.processing_class = tokenizer
-        print(f"[OK] Model loaded\n")
-
-    # ── Standard Loading (no quantization) ─────────────────────────────────────
-    else:
-        tokenizer = AutoTokenizer.from_pretrained(args.model)
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        model = AutoModelForCausalLM.from_pretrained(
-            args.model,
-            dtype=torch.float16 if device == "cuda" else torch.float32,
-            device_map="auto" if device == "cuda" else None,
-        )
-        if device == "cpu":
-            model = model.to(device)
-        if not hasattr(model, "processing_class"):
-            model.processing_class = tokenizer
-        print(f"[OK] Model loaded\n")
-
-    # ── Training Loop ─────────────────────────────────────────────────────────
-
-    reward_history: dict[str, list[float]] = {t: [] for t in tasks}
-    all_trajectories: list[dict] = []
-
-    # Checkpoint dir
-    CHECKPOINT_DIR = "./phase2_checkpoints"
-    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
-
-    for task_id in tasks:
-        print(f"\n{'='*60}")
-        print(f"[TRAIN] Training on task: {task_id}")
-        print(f"{'='*60}")
-
-        task_rewards = []
-
-        for ep in range(1, args.episodes + 1):
-            seed = ep  # different seed each episode = different incident
-
-            total_reward, steps, trajectory = run_episode(
-                env=env,
-                model=model,
-                tokenizer=tokenizer,
-                task_id=task_id,
-                seed=seed,
-                device=device,
-                verbose=args.verbose,
-            )
-
-            task_rewards.append(total_reward)
-            all_trajectories.extend(trajectory)
-
-            # Rolling average for display
-            window = task_rewards[-10:]
-            rolling_avg = sum(window) / len(window)
-
-            # Save checkpoint every 25 episodes
-            if ep % 25 == 0:
-                ckpt_path = os.path.join(CHECKPOINT_DIR, f"{task_id}_ep{ep}.json")
-                with open(ckpt_path, "w") as f:
-                    json.dump({
-                        "task_id": task_id,
-                        "episode": ep,
-                        "rewards": task_rewards,
-                    }, f)
-                print(f"  [CHECKPOINT] Saved {task_id} ep{ep} -> {ckpt_path}")
-
-            print(
-                f"  Episode {ep:3d}/{args.episodes} | "
-                f"Reward: {total_reward:+.3f} | "
-                f"Steps: {steps:2d} | "
-                f"Rolling avg (10): {rolling_avg:.3f}"
-            )
-
-            # Small delay to avoid hammering the env
-            time.sleep(0.1)
-
-        reward_history[task_id] = task_rewards
-
-        # Summary for this task
-        if task_rewards:
-            first_10 = sum(task_rewards[:10]) / min(10, len(task_rewards))
-            last_10 = sum(task_rewards[-10:]) / min(10, len(task_rewards))
-            improvement = last_10 - first_10
-            print(f"\n[STATS] {task_id} Summary:")
-            print(f"     First 10 episodes avg: {first_10:.3f}")
-            print(f"     Last  10 episodes avg: {last_10:.3f}")
-            print(f"     Improvement:           {improvement:+.3f}")
-
-    # ── Save Reward Curve ─────────────────────────────────────────────────────
-
-    save_reward_curve(reward_history, "reward_curve.png")
-
-    # ── GRPO Fine-tuning Pass ─────────────────────────────────────────────────
-    if all_trajectories:
-        print(f"\n[GRPO] Collected {len(all_trajectories)} trajectory steps from rollout.")
-
-        if args.skip_grpo:
-            print("[GRPO] Skipping GRPO fine-tuning (--skip_grpo set).")
-            print("[GRPO] Reward curves from rollout demonstrate training progress.")
-        else:
-            # Reward is carried from the rollout trajectory and fed into GRPO as a verifiable scalar.
-            def reward_fn(completions, **kwargs):
-                rewards = kwargs.get("reward", None)
-                if rewards is None:
-                    return [0.0 for _ in completions]
-                return [float(r) for r in rewards]
-
-            try:
-                grpo_dataset = build_grpo_dataset(all_trajectories)
-                max_steps = min(max(1, args.grpo_max_steps), max(1, len(grpo_dataset)))
-
-                print(f"[GRPO] Running GRPO fine-tuning on {len(grpo_dataset)} trajectory steps...")
-
-                # Keep memory pressure low for Colab T4 / laptop GPUs.
-                if hasattr(model, "config"):
-                    model.config.use_cache = False
-
-                grpo_args = GRPOConfig(
-                    output_dir=args.output_dir,
-                    per_device_train_batch_size=1,
-                    gradient_accumulation_steps=4,
-                    num_train_epochs=1,
-                    max_steps=max_steps,
-                    learning_rate=1e-5,
-                    logging_steps=10,
-                    save_steps=100,
-                    report_to=[],
-                )
-
-                trainer = GRPOTrainer(
-                    model=model,
-                    reward_funcs=reward_fn,
-                    args=grpo_args,
-                    train_dataset=grpo_dataset,
-                    processing_class=tokenizer,
-                )
-
-                train_output = trainer.train()
-                metrics = getattr(train_output, "metrics", None)
-                if metrics:
-                    print(f"[GRPO] Metrics: {metrics}")
-                print("[OK] GRPO training complete")
-
-            except RuntimeError as e:
-                if "out of memory" in str(e).lower():
-                    print(f"[WARN] GRPO OOM: {e}")
-                    print("[WARN] Continuing with rollout-only results. Try --skip_grpo or lower --grpo_max_steps.")
-                else:
-                    raise
-            except Exception as e:
-                print(f"[WARN] GRPO trainer error: {e}")
-                print("[WARN] Continuing with rollout-only results.")
-
-    # ── Save Model ────────────────────────────────────────────────────────────
-
-    os.makedirs(args.output_dir, exist_ok=True)
-    # Clear CUDA state and move to CPU before saving
-    try:
-        if device == "cuda":
-            torch.cuda.empty_cache()
-    except Exception:
-        pass
-
-    # Merge LoRA adapter before saving (for LoRA models)
-    if use_lora and hasattr(model, "merge_and_unload"):
-        print("[SAVE] Merging LoRA adapter into base weights...")
-        model = model.merge_and_unload()
-        print("[OK] LoRA merged — saving full model")
-    elif use_unsloth:
-        print("[SAVE] Unsloth model — saving merged weights")
-    elif getattr(args, "load_in_4bit", False):
-        print("[SAVE] BitsAndBytes QLoRA model — saving adapter")
-
-    model = model.cpu()
-    model.save_pretrained(args.output_dir)
-    tokenizer.save_pretrained(args.output_dir)
-    print(f"\n[SAVE] Model saved -> {args.output_dir}")
-
-    # ── Push to Hub ───────────────────────────────────────────────────────────
-
-    if args.push_to_hub and args.hub_model_id:
-        print(f"\n[PUSH] Pushing to HuggingFace Hub: {args.hub_model_id}")
-        model.push_to_hub(args.hub_model_id)
-        tokenizer.push_to_hub(args.hub_model_id)
-        print(f"[OK] Model pushed -> https://huggingface.co/{args.hub_model_id}")
-
-    # ── Final Summary ─────────────────────────────────────────────────────────
-
-    print(f"\n{'='*60}")
-    print(f"[OK] TRAINING COMPLETE")
-    print(f"{'='*60}")
-    print(f"  Reward curve:  reward_curve.png")
-    print(f"  Trained model: {args.output_dir}")
-    if args.push_to_hub and args.hub_model_id:
-        print(f"  HF Hub:        https://huggingface.co/{args.hub_model_id}")
-    print(f"\n  Use reward_curve.png in your demo slide.")
-    print(f"  This image is 20% of your judging score.\n")
-
-
-if __name__ == "__main__":
-    main()
+"""
+train.py — LogTriageEnv GRPO Training Loop
+Meta × PyTorch × Scaler OpenEnv Hackathon — Grand Finale
+
+Usage:
+    python train.py --model HuggingFaceTB/SmolLM2-360M-Instruct --task single_crash --episodes 50 --env_url http://localhost:7860
+    python train.py --model HuggingFaceTB/SmolLM2-360M-Instruct --task all --episodes 100 --env_url http://localhost:7860
+
+    # Colab T4 GPU — use Unsloth (recommended for Qwen 3B/7B):
+    python train.py --model Qwen/Qwen2.5-7B-Instruct --task all --episodes 50 --use_unsloth --env_url https://ogrohit-logtriage-env.hf.space
+    python train.py --model Qwen/Qwen2.5-3B-Instruct --task all --episodes 50 --use_unsloth --env_url https://ogrohit-logtriage-env.hf.space
+
+    # Local laptop (no quantization):
+    python train.py --model HuggingFaceTB/SmolLM2-360M-Instruct --task all --episodes 50 --env_url http://localhost:7860
+
+    # Onsite with A100 — use Unsloth for max speed:
+    python train.py --model Qwen/Qwen2.5-32B-Instruct --task all --episodes 100 --use_unsloth --env_url https://ogrohit-logtriage-env.hf.space
+"""
+
+import argparse
+import json
+import re
+import time
+import os
+from dataclasses import dataclass, field
+from typing import Optional, List
+
+import requests
+import matplotlib.pyplot as plt
+import matplotlib
+matplotlib.use("Agg")  # headless — no display required
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from trl import GRPOConfig, GRPOTrainer
+from datasets import Dataset
+
+try:
+    from peft import LoraConfig, get_peft_model, PeftModel
+    PEFT_AVAILABLE = True
+except ImportError:
+    PEFT_AVAILABLE = False
+
+try:
+    from unsloth import FastLanguageModel
+    UNSLOTH_AVAILABLE = True
+except Exception:
+    UNSLOTH_AVAILABLE = False
+
+# ── Constants ────────────────────────────────────────────────────────────────
+
+VALID_ACTION_TYPES = [
+    "classify_severity",
+    "identify_root_cause",
+    "escalate",
+    "remediate",
+    "request_more_logs",
+    "resolve",
+    "ignore",
+]
+
+VALID_VALUES = {
+    "classify_severity": ["P1", "P2", "P3"],
+    "identify_root_cause": [
+        "api-gateway", "auth-service", "user-db",
+        "payment-service", "payment-db",
+        "notification-service", "email-queue",
+    ],
+    "escalate": ["sre-team", "backend-team", "dba-team", "security-team", "ignore"],
+    "remediate": [
+        "restart:api-gateway", "restart:auth-service", "restart:user-db",
+        "restart:payment-service", "restart:payment-db",
+        "restart:notification-service", "restart:email-queue",
+        "rollback:api-gateway", "rollback:auth-service", "rollback:payment-service",
+        "scale:api-gateway", "scale:payment-service",
+        "flush-cache:user-db", "flush-cache:payment-db",
+        "kill-query:user-db", "kill-query:payment-db",
+    ],
+    "request_more_logs": [
+        "api-gateway", "auth-service", "user-db",
+        "payment-service", "payment-db",
+        "notification-service", "email-queue", "all",
+    ],
+    "resolve": ["resolved"],
+    "ignore": ["noise"],
+}
+
+SYSTEM_PROMPT = """You are an expert SRE (Site Reliability Engineer) triaging a live production incident.
+
+You will receive log lines from a microservice cluster. Your job is to reason carefully and take ONE action per step.
+
+The service topology is:
+  [api-gateway] → [auth-service] → [user-db]
+               → [payment-service] → [payment-db]
+               → [notification-service] → [email-queue]
+
+Available actions:
+- classify_severity: Set priority. Values: P1 (customer-facing outage), P2 (degradation), P3 (warning)
+- identify_root_cause: Point to the failing service. Values: api-gateway, auth-service, user-db, payment-service, payment-db, notification-service, email-queue
+- escalate: Page a team. Values: sre-team, backend-team, dba-team, security-team, ignore
+- remediate: Apply a fix. Values: restart:<service>, rollback:<service>, scale:<service>, flush-cache:<service>, kill-query:<service>
+- request_more_logs: Get more logs. Values: <service-name> or all
+- resolve: Mark incident resolved. Value: resolved
+- ignore: Mark as noise. Value: noise
+
+CRITICAL RULES:
+1. For cascading failures, find the ROOT CAUSE service, not the first service that shows errors
+2. P1 = customer-facing impact (error rate >5%), P2 = degradation, P3 = warning only
+3. Do NOT over-escalate. Paging the wrong team is penalized.
+4. Be efficient — unnecessary steps reduce your score.
+
+You MUST respond in this exact JSON format and nothing else:
+{
+  "action_type": "<one of the action types above>",
+  "value": "<valid value for that action type>",
+  "confidence": <float 0.0-1.0>,
+  "reasoning": "<one sentence explaining why>"
+}"""
+
+
+# ── Env Client ───────────────────────────────────────────────────────────────
+
+class LogTriageEnvClient:
+    """HTTP client for LogTriageEnv."""
+
+    def __init__(self, base_url: str):
+        self.base_url = base_url.rstrip("/")
+        self._verify_connection()
+
+    def _verify_connection(self):
+        try:
+            r = requests.get(f"{self.base_url}/health", timeout=10)
+            r.raise_for_status()
+            print(f"[OK] Connected to LogTriageEnv at {self.base_url}")
+        except Exception as e:
+            raise RuntimeError(
+                f"[ERROR] Cannot reach LogTriageEnv at {self.base_url}\n"
+                f"   Make sure Docker is running: docker run -p 7860:7860 logtriage-env\n"
+                f"   Error: {e}"
+            )
+
+    def reset(self, task_id: str, seed: int = 42) -> dict:
+        r = requests.post(
+            f"{self.base_url}/reset",
+            json={"task_id": task_id, "seed": seed},
+            timeout=15,
+        )
+        r.raise_for_status()
+        return r.json()
+
+    def step(self, action: dict) -> dict:
+        r = requests.post(
+            f"{self.base_url}/step",
+            json=action,
+            timeout=15,
+        )
+        r.raise_for_status()
+        return r.json()
+
+    def get_tasks(self) -> list:
+        r = requests.get(f"{self.base_url}/tasks", timeout=10)
+        r.raise_for_status()
+        return r.json()["tasks"]
+
+
+# ── Observation Formatting ───────────────────────────────────────────────────
+
+def format_observation(obs: dict, step: int) -> str:
+    """Convert raw env observation dict into a clean prompt string."""
+    lines = []
+
+    lines.append(f"=== INCIDENT TRIAGE — Step {step} ===")
+    lines.append(f"Incident ID: {obs.get('incident_id', 'unknown')}")
+    lines.append(f"Active Alerts: {', '.join(obs.get('active_alerts', []))}")
+    lines.append("")
+
+    # System state
+    lines.append("--- System State ---")
+    system_state = obs.get("system_state", {})
+    for svc, status in system_state.items():
+        if isinstance(status, dict):
+            lines.append(
+                f"  {svc}: {status.get('status','?')} | "
+                f"error_rate={status.get('error_rate', 0):.1%} | "
+                f"p99={status.get('latency_p99_ms', 0)}ms"
+            )
+        else:
+            lines.append(f"  {svc}: {status}")
+
+    # Log lines
+    lines.append("")
+    lines.append("--- Log Stream ---")
+    logs = obs.get("logs", [])
+    if isinstance(logs, list):
+        for log in logs[-15:]:  # last 15 lines to stay within context
+            if isinstance(log, dict):
+                ts = log.get("timestamp", "")
+                level = log.get("level", "")
+                svc = log.get("service", "")
+                msg = log.get("message", "")
+                lines.append(f"  [{ts}] {level:5} {svc:25} {msg}")
+            else:
+                lines.append(f"  {log}")
+    else:
+        lines.append(str(logs))
+
+    # Feedback from last action
+    feedback = obs.get("last_action_feedback", "")
+    if feedback:
+        lines.append("")
+        lines.append(f"--- Last Action Feedback ---")
+        lines.append(f"  {feedback}")
+
+    lines.append("")
+    lines.append("What is your next action? Respond in JSON only.")
+
+    return "\n".join(lines)
+
+
+# ── Action Parsing ────────────────────────────────────────────────────────────
+
+def parse_action(llm_output: str) -> Optional[dict]:
+    """
+    Parse LLM output into a valid TriageAction dict.
+    Returns None if parsing fails completely.
+    """
+    # Try direct JSON parse first
+    try:
+        # Strip markdown code fences if present
+        clean = re.sub(r"```(?:json)?", "", llm_output).strip().rstrip("```").strip()
+        # Find first { ... } block
+        match = re.search(r"\{.*\}", clean, re.DOTALL)
+        if match:
+            action = json.loads(match.group())
+            if "action_type" in action and "value" in action:
+                # Validate action_type
+                if action["action_type"] not in VALID_ACTION_TYPES:
+                    return None
+                # Validate value against strict server-side rules
+                validated = _validate_action_value(action["action_type"], action.get("value", ""))
+                if validated is None:
+                    return None
+                action["value"] = validated
+                action["confidence"] = 0.5
+                action["reasoning"] = ""
+                return action
+    except (json.JSONDecodeError, KeyError):
+        pass
+
+    # Fallback: keyword extraction (only on known-good pairs)
+    output_lower = llm_output.lower()
+    for action_type in VALID_ACTION_TYPES:
+        if action_type.replace("_", " ") in output_lower or action_type in output_lower:
+            for value in VALID_VALUES.get(action_type, []):
+                if value.lower() in output_lower:
+                    # Extra validation for escalate: "ignore" is NOT a valid escalate value
+                    if action_type == "escalate" and value == "ignore":
+                        continue
+                    return {
+                        "action_type": action_type,
+                        "value": value,
+                        "confidence": 0.3,
+                        "reasoning": "parsed via fallback",
+                    }
+
+    # Last resort: safe default
+    return {
+        "action_type": "request_more_logs",
+        "value": "all",
+        "confidence": 0.1,
+        "reasoning": "failed to parse LLM output",
+    }
+
+
+def _validate_action_value(action_type: str, value: str) -> Optional[str]:
+    """Validate action value against server-side rules. Returns clean value or None."""
+    if action_type == "classify_severity":
+        if value in ("P1", "P2", "P3"):
+            return value
+    elif action_type == "identify_root_cause":
+        valid = {
+            "api-gateway", "auth-service", "user-db",
+            "payment-service", "payment-db",
+            "notification-service", "email-queue",
+        }
+        if value in valid:
+            return value
+        # Fuzzy match: "payment" -> "payment-service"
+        if value in ("payment", "payment svc", "paymentservice"):
+            return "payment-service"
+        if value in ("user", "userdb", "user_db"):
+            return "user-db"
+        if value in ("auth", "authsvc"):
+            return "auth-service"
+        if value in ("api", "gateway", "api-gw"):
+            return "api-gateway"
+        if value in ("notif", "notification", "notif-service"):
+            return "notification-service"
+        if value in ("email", "emailqueue", "queue"):
+            return "email-queue"
+    elif action_type == "escalate":
+        valid = {"sre-team", "backend-team", "dba-team", "security-team"}
+        if value in valid:
+            return value
+    elif action_type == "remediate":
+        if ":" in value:
+            prefix, service = value.split(":", 1)
+            valid_prefixes = {"restart", "rollback", "scale", "flush-cache", "kill-query"}
+            if prefix in valid_prefixes:
+                # Map service aliases
+                service_map = {
+                    "payment": "payment-service",
+                    "userdb": "user-db",
+                    "user_db": "user-db",
+                    "auth": "auth-service",
+                    "api": "api-gateway",
+                    "gateway": "api-gateway",
+                    "notif": "notification-service",
+                    "email": "email-queue",
+                }
+                clean_service = service_map.get(service, service)
+                return f"{prefix}:{clean_service}"
+    elif action_type == "request_more_logs":
+        valid_services = {
+            "api-gateway", "auth-service", "user-db",
+            "payment-service", "payment-db",
+            "notification-service", "email-queue", "all",
+        }
+        if value in valid_services:
+            return value
+        service_map = {
+            "payment": "payment-service", "userdb": "user-db",
+            "user_db": "user-db", "auth": "auth-service",
+            "api": "api-gateway", "gateway": "api-gateway",
+            "notif": "notification-service", "email": "email-queue",
+        }
+        if value in service_map:
+            return service_map[value]
+    elif action_type == "resolve":
+        if value == "resolved":
+            return "resolved"
+    elif action_type == "ignore":
+        if value == "noise":
+            return "noise"
+    return None
+
+
+# ── Single Episode Rollout ───────────────────────────────────────────────────
+
+def run_episode(
+    env: LogTriageEnvClient,
+    model,
+    tokenizer,
+    task_id: str,
+    seed: int,
+    device: str,
+    max_steps: int = 15,
+    verbose: bool = False,
+) -> tuple[float, int, list[dict]]:
+    """
+    Run one full episode.
+    Returns: (total_reward, steps_taken, trajectory)
+    trajectory = list of {prompt, response, reward} dicts for GRPO
+    """
+    obs = env.reset(task_id=task_id, seed=seed)
+    total_reward = 0.0
+    steps = 0
+    trajectory = []
+    done = False
+
+    while not done and steps < max_steps:
+        # Format observation into prompt
+        prompt_text = format_observation(obs, steps + 1)
+
+        # Build chat messages
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": prompt_text},
+        ]
+
+        # Tokenize
+        input_ids = tokenizer.apply_chat_template(
+            messages,
+            return_tensors="pt",
+            add_generation_prompt=True,
+        )
+        # HF tokenizers may return a tensor directly or a BatchEncoding.
+        if isinstance(input_ids, torch.Tensor):
+            input_ids = input_ids.to(device)
+        else:
+            input_ids = input_ids["input_ids"].to(device)
+        pad_token_id = tokenizer.pad_token_id
+        if pad_token_id is None:
+            pad_token_id = tokenizer.eos_token_id
+        attention_mask = (input_ids != pad_token_id).long()
+        gen_kwargs = {
+            "max_new_tokens": 150,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "attention_mask": attention_mask,
+            "pad_token_id": tokenizer.eos_token_id,
+        }
+
+        # Generate
+        with torch.no_grad():
+            output_ids = model.generate(input_ids, **gen_kwargs)
+
+        # Decode only the new tokens
+        prompt_len = input_ids.shape[1]
+        new_tokens = output_ids[0][prompt_len:]
+        llm_output = tokenizer.decode(new_tokens, skip_special_tokens=True)
+
+        # Parse action
+        action = parse_action(llm_output)
+        if action is None:
+            action = {"action_type": "request_more_logs", "value": "all",
+                      "confidence": 0.1, "reasoning": "parse failed"}
+
+        # Step env
+        try:
+            obs = env.step(action)
+        except requests.HTTPError as e:
+            if verbose:
+                print(f"[WARN] Step HTTP error: {e}")
+            break
+
+        # Extract reward
+        step_reward = obs.get("reward", 0.0)
+        total_reward += step_reward
+        done = obs.get("done", False)
+        steps += 1
+
+        # Store for GRPO
+        trajectory.append({
+            "prompt": prompt_text,
+            "response": llm_output,
+            "reward": step_reward,
+        })
+
+        if verbose:
+            print(f"    Step {steps}: action={action['action_type']}({action['value']}) "
+                  f"reward={step_reward:+.2f} done={done}")
+
+    return total_reward, steps, trajectory
+
+
+# ── Reward Curve Plot ─────────────────────────────────────────────────────────
+
+def save_reward_curve(history: dict[str, list[float]], output_path: str = "reward_curve.png"):
+    """
+    history: {"single_crash": [r1, r2, ...], "cascading_failure": [...], ...}
+    """
+    fig, ax = plt.subplots(figsize=(10, 6))
+
+    colors = {"single_crash": "#00C49F", "cascading_failure": "#FFBB28", "silent_degradation": "#FF6B6B"}
+    labels = {"single_crash": "Task 1: Single Crash (Easy)",
+              "cascading_failure": "Task 2: Cascading Failure (Medium)",
+              "silent_degradation": "Task 3: Silent Degradation (Hard)"}
+
+    for task_id, rewards in history.items():
+        if not rewards:
+            continue
+        # Smooth with rolling average (window=5)
+        smoothed = []
+        for i in range(len(rewards)):
+            window = rewards[max(0, i-4):i+1]
+            smoothed.append(sum(window) / len(window))
+
+        episodes = list(range(1, len(rewards) + 1))
+        color = colors.get(task_id, "#8884d8")
+        label = labels.get(task_id, task_id)
+
+        ax.plot(episodes, rewards, alpha=0.3, color=color, linewidth=0.8)
+        ax.plot(episodes, smoothed, color=color, linewidth=2.5, label=label)
+
+    ax.set_xlabel("Episode", fontsize=13)
+    ax.set_ylabel("Episode Reward", fontsize=13)
+    ax.set_title("LogTriageEnv — Agent Reward Improvement During GRPO Training", fontsize=14, fontweight="bold")
+    ax.legend(fontsize=11)
+    ax.grid(True, alpha=0.3)
+    ax.set_ylim(bottom=0)
+
+    # Add annotation
+    ax.annotate(
+        "Higher = agent solves incidents faster with fewer wrong actions",
+        xy=(0.02, 0.02), xycoords="axes fraction",
+        fontsize=9, color="gray", style="italic"
+    )
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"[PLOT] Reward curve saved -> {output_path}")
+
+
+# ── GRPO Dataset Builder ──────────────────────────────────────────────────────
+
+def build_grpo_dataset(trajectories: list[dict]) -> Dataset:
+    """
+    Build a HF Dataset from collected trajectories for GRPOTrainer.
+    Format: {"prompt": str, "completion": str, "reward": float}
+    """
+    if not trajectories:
+        # Return minimal dummy dataset if no trajectories yet
+        return Dataset.from_dict({
+            "prompt": ["dummy"],
+            "completion": ["{}"],
+            "reward": [0.0],
+        })
+
+    return Dataset.from_dict({
+        "prompt": [t["prompt"] for t in trajectories],
+        "completion": [t["response"] for t in trajectories],
+        "reward": [t["reward"] for t in trajectories],
+    })
+
+
+# ── Main Training Loop ────────────────────────────────────────────────────────
+
+def main():
+    parser = argparse.ArgumentParser(description="LogTriageEnv GRPO Training")
+    parser.add_argument("--model", default="HuggingFaceTB/SmolLM2-360M-Instruct",
+                        help="HuggingFace model ID")
+    parser.add_argument("--task", default="single_crash",
+                        choices=["single_crash", "cascading_failure", "silent_degradation", "all"],
+                        help="Task to train on. 'all' trains on all 3.")
+    parser.add_argument("--episodes", type=int, default=50,
+                        help="Number of training episodes per task")
+    parser.add_argument("--env_url", default="http://localhost:7860",
+                        help="LogTriageEnv base URL")
+    parser.add_argument("--output_dir", default="./logtriage-trained",
+                        help="Where to save the trained model")
+    parser.add_argument("--push_to_hub", action="store_true",
+                        help="Push trained model to HuggingFace Hub")
+    parser.add_argument("--hub_model_id", default=None,
+                        help="HF Hub model ID (e.g. username/logtriage-sre-agent)")
+    parser.add_argument("--verbose", action="store_true",
+                        help="Print step-by-step actions during episodes")
+    parser.add_argument("--load_in_4bit", action="store_true",
+                        help="Load model with 4-bit QLoRA quantization via BitsAndBytes (for large models on limited VRAM)")
+    parser.add_argument("--use_unsloth", action="store_true",
+                        help="Load model using Unsloth (recommended for Qwen on T4/A100 — faster and more memory efficient)")
+    parser.add_argument("--skip_grpo", action="store_true",
+                        help="Skip GRPO fine-tuning and only run rollout episodes (useful when debugging or avoiding OOM)")
+    parser.add_argument("--grpo_max_steps", type=int, default=35,
+                        help="Maximum GRPO optimization steps after rollout (default: 35)")
+    args = parser.parse_args()
+
+    # ── Setup ────────────────────────────────────────────────────────────────
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print("\n[LOGGING] LogTriageEnv GRPO Training")
+    print(f"   Model:   {args.model}")
+    print(f"   Task:    {args.task}")
+    print(f"   Episodes: {args.episodes}")
+    print(f"   Device:  {device}")
+    print(f"   Env URL: {args.env_url}\n")
+
+    # Connect to env
+    env = LogTriageEnvClient(args.env_url)
+
+    # Determine tasks to train on
+    if args.task == "all":
+        tasks = ["single_crash", "cascading_failure", "silent_degradation"]
+    else:
+        tasks = [args.task]
+
+    # Load model + tokenizer
+    print(f"[MODEL] Loading model: {args.model}")
+    use_unsloth = getattr(args, "use_unsloth", False)
+    use_lora = False
+
+    # ── Unsloth Path (recommended for Qwen on T4/A100) ───────────────────────
+    if use_unsloth and device == "cuda" and UNSLOTH_AVAILABLE:
+        print("[UNSLOTH] Loading model with Unsloth...")
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=args.model,
+            max_seq_length=2048,
+            load_in_4bit=True,
+            dtype=None,  # Auto-detect
+        )
+        print(f"[OK] Model loaded via Unsloth (4-bit)")
+
+        # Apply LoRA via Unsloth
+        print("[UNSLOTH] Applying LoRA adapter (r=16, alpha=32)...")
+        model = FastLanguageModel.get_peft_model(
+            model,
+            r=16,
+            lora_alpha=32,
+            target_modules=[
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj",
+            ],
+            lora_dropout=0.05,
+            bias="none",
+        )
+        model.print_trainable_parameters()
+        use_lora = True
+        print(f"[OK] Unsloth LoRA attached")
+        print(f"[OK] Model loaded\n")
+
+    # ── BitsAndBytes QLoRA Path (manual, or fallback) ─────────────────────────
+    elif getattr(args, "load_in_4bit", False) and device == "cuda":
+        print("[QLoRA] Loading model with BitsAndBytes 4-bit...")
+        tokenizer = AutoTokenizer.from_pretrained(args.model)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+        )
+        print(f"[OK] 4-bit BitsAndBytesConfig applied")
+
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model,
+            quantization_config=bnb_config,
+            device_map="auto",
+        )
+        print(f"[OK] Model loaded in 4-bit quantized mode")
+
+        if PEFT_AVAILABLE:
+            print("[QLoRA] Applying LoRA adapter...")
+            lora_config = LoraConfig(
+                r=16,
+                lora_alpha=32,
+                target_modules=[
+                    "q_proj", "k_proj", "v_proj", "o_proj",
+                    "gate_proj", "up_proj", "down_proj",
+                ],
+                lora_dropout=0.05,
+                bias="none",
+                task_type="CAUSAL_LM",
+            )
+            model = get_peft_model(model, lora_config)
+            model.print_trainable_parameters()
+            use_lora = True
+            print(f"[OK] LoRA adapter attached (r=16, alpha=32)")
+        else:
+            print("[WARN] PEFT not installed. Using quantized model without LoRA.")
+
+        if not hasattr(model, "processing_class"):
+            model.processing_class = tokenizer
+        print(f"[OK] Model loaded\n")
+
+    # ── Standard Loading (no quantization) ─────────────────────────────────────
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(args.model)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model,
+            dtype=torch.float16 if device == "cuda" else torch.float32,
+            device_map="auto" if device == "cuda" else None,
+        )
+        if device == "cpu":
+            model = model.to(device)
+        if not hasattr(model, "processing_class"):
+            model.processing_class = tokenizer
+        print(f"[OK] Model loaded\n")
+
+    # ── Training Loop ─────────────────────────────────────────────────────────
+
+    reward_history: dict[str, list[float]] = {t: [] for t in tasks}
+    all_trajectories: list[dict] = []
+
+    # Checkpoint dir
+    CHECKPOINT_DIR = "./phase2_checkpoints"
+    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+
+    for task_id in tasks:
+        print(f"\n{'='*60}")
+        print(f"[TRAIN] Training on task: {task_id}")
+        print(f"{'='*60}")
+
+        task_rewards = []
+
+        for ep in range(1, args.episodes + 1):
+            seed = ep  # different seed each episode = different incident
+
+            total_reward, steps, trajectory = run_episode(
+                env=env,
+                model=model,
+                tokenizer=tokenizer,
+                task_id=task_id,
+                seed=seed,
+                device=device,
+                verbose=args.verbose,
+            )
+
+            task_rewards.append(total_reward)
+            all_trajectories.extend(trajectory)
+
+            # Rolling average for display
+            window = task_rewards[-10:]
+            rolling_avg = sum(window) / len(window)
+
+            # Save checkpoint every 25 episodes
+            if ep % 25 == 0:
+                ckpt_path = os.path.join(CHECKPOINT_DIR, f"{task_id}_ep{ep}.json")
+                with open(ckpt_path, "w") as f:
+                    json.dump({
+                        "task_id": task_id,
+                        "episode": ep,
+                        "rewards": task_rewards,
+                    }, f)
+                print(f"  [CHECKPOINT] Saved {task_id} ep{ep} -> {ckpt_path}")
+
+            print(
+                f"  Episode {ep:3d}/{args.episodes} | "
+                f"Reward: {total_reward:+.3f} | "
+                f"Steps: {steps:2d} | "
+                f"Rolling avg (10): {rolling_avg:.3f}"
+            )
+
+            # Small delay to avoid hammering the env
+            time.sleep(0.1)
+
+        reward_history[task_id] = task_rewards
+
+        # Summary for this task
+        if task_rewards:
+            first_10 = sum(task_rewards[:10]) / min(10, len(task_rewards))
+            last_10 = sum(task_rewards[-10:]) / min(10, len(task_rewards))
+            improvement = last_10 - first_10
+            print(f"\n[STATS] {task_id} Summary:")
+            print(f"     First 10 episodes avg: {first_10:.3f}")
+            print(f"     Last  10 episodes avg: {last_10:.3f}")
+            print(f"     Improvement:           {improvement:+.3f}")
+
+    # ── Save Reward Curve ─────────────────────────────────────────────────────
+
+    save_reward_curve(reward_history, "reward_curve.png")
+
+    # ── GRPO Fine-tuning Pass ─────────────────────────────────────────────────
+    if all_trajectories:
+        print(f"\n[GRPO] Collected {len(all_trajectories)} trajectory steps from rollout.")
+
+        if args.skip_grpo:
+            print("[GRPO] Skipping GRPO fine-tuning (--skip_grpo set).")
+            print("[GRPO] Reward curves from rollout demonstrate training progress.")
+        else:
+            # Reward is carried from the rollout trajectory and fed into GRPO as a verifiable scalar.
+            def reward_fn(completions, **kwargs):
+                rewards = kwargs.get("reward", None)
+                if rewards is None:
+                    return [0.0 for _ in completions]
+                return [float(r) for r in rewards]
+
+            try:
+                grpo_dataset = build_grpo_dataset(all_trajectories)
+                max_steps = min(max(1, args.grpo_max_steps), max(1, len(grpo_dataset)))
+
+                print(f"[GRPO] Running GRPO fine-tuning on {len(grpo_dataset)} trajectory steps...")
+
+                # Keep memory pressure low for Colab T4 / laptop GPUs.
+                if hasattr(model, "config"):
+                    model.config.use_cache = False
+
+                use_bf16 = device == "cuda" and torch.cuda.is_bf16_supported()
+                use_fp16 = device == "cuda" and not use_bf16
+                if use_bf16:
+                    print("[GRPO] Precision: bf16")
+                elif use_fp16:
+                    print("[GRPO] Precision: fp16 (bf16 unsupported on this GPU)")
+                else:
+                    print("[GRPO] Precision: fp32 (CPU mode)")
+
+                grpo_args = GRPOConfig(
+                    output_dir=args.output_dir,
+                    per_device_train_batch_size=1,
+                    gradient_accumulation_steps=4,
+                    num_train_epochs=1,
+                    max_steps=max_steps,
+                    learning_rate=1e-5,
+                    generation_batch_size=4,
+                    num_generations=4,
+                    logging_steps=10,
+                    save_steps=100,
+                    report_to=[],
+                    bf16=use_bf16,
+                    fp16=use_fp16,
+                )
+
+                trainer = GRPOTrainer(
+                    model=model,
+                    reward_funcs=reward_fn,
+                    args=grpo_args,
+                    train_dataset=grpo_dataset,
+                    processing_class=tokenizer,
+                )
+
+                train_output = trainer.train()
+                metrics = getattr(train_output, "metrics", None)
+                if metrics:
+                    print(f"[GRPO] Metrics: {metrics}")
+                print("[OK] GRPO training complete")
+
+            except RuntimeError as e:
+                if "out of memory" in str(e).lower():
+                    print(f"[WARN] GRPO OOM: {e}")
+                    print("[WARN] Continuing with rollout-only results. Try --skip_grpo or lower --grpo_max_steps.")
+                else:
+                    raise
+            except Exception as e:
+                print(f"[WARN] GRPO trainer error: {e}")
+                print("[WARN] Continuing with rollout-only results.")
+
+    # ── Save Model ────────────────────────────────────────────────────────────
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    # Clear CUDA state and move to CPU before saving
+    try:
+        if device == "cuda":
+            torch.cuda.empty_cache()
+    except Exception:
+        pass
+
+    # Merge LoRA adapter before saving (for LoRA models)
+    if use_lora and hasattr(model, "merge_and_unload"):
+        print("[SAVE] Merging LoRA adapter into base weights...")
+        model = model.merge_and_unload()
+        print("[OK] LoRA merged — saving full model")
+    elif use_unsloth:
+        print("[SAVE] Unsloth model — saving merged weights")
+    elif getattr(args, "load_in_4bit", False):
+        print("[SAVE] BitsAndBytes QLoRA model — saving adapter")
+
+    model = model.cpu()
+    model.save_pretrained(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+    print(f"\n[SAVE] Model saved -> {args.output_dir}")
+
+    # ── Push to Hub ───────────────────────────────────────────────────────────
+
+    if args.push_to_hub and args.hub_model_id:
+        print(f"\n[PUSH] Pushing to HuggingFace Hub: {args.hub_model_id}")
+        model.push_to_hub(args.hub_model_id)
+        tokenizer.push_to_hub(args.hub_model_id)
+        print(f"[OK] Model pushed -> https://huggingface.co/{args.hub_model_id}")
+
+    # ── Final Summary ─────────────────────────────────────────────────────────
+
+    print(f"\n{'='*60}")
+    print(f"[OK] TRAINING COMPLETE")
+    print(f"{'='*60}")
+    print(f"  Reward curve:  reward_curve.png")
+    print(f"  Trained model: {args.output_dir}")
+    if args.push_to_hub and args.hub_model_id:
+        print(f"  HF Hub:        https://huggingface.co/{args.hub_model_id}")
+    print(f"\n  Use reward_curve.png in your demo slide.")
+    print(f"  This image is 20% of your judging score.\n")
+
+
+if __name__ == "__main__":
+    main()