train.py

"""
Training script for the DevOps Incident Triage Environment.
Uses Unsloth + TRL for a two-phase training pipeline:

Phase 1 — SFT Warm-start (NEW)
  1 epoch of supervised fine-tuning on the optimal action sequences from all
  6 task scenarios.  This gives the model a non-random starting policy so that
  GRPO receives non-zero rewards from step 1.

Phase 2 — GRPO RL Training
  Curriculum design (all 6 tasks):
    - 15 prompts per task initial state  (easy/medium/hard/network/memory_leak/disk_full)
    - 15 mid-episode prompts (1-2 random steps taken first across all tasks)
  Shuffled before training.  max_steps bumped to 200 for a more visible reward curve.

Reward functions:
  1. format_reward_func         — valid action string check
  2. step_reward_func           — environment step reward (task-aware)
  3. anti_cheat_reward_func     — penalise no_op / premature resolve
  4. task_alignment_reward_func — rewards diagnostics/mitigations matching the task

Note: torch / unsloth / trl are imported lazily inside main() so that the
reward functions and dataset helpers can be imported on machines without a GPU.
"""

import csv
import json
import os
import random
from typing import List

from datasets import Dataset

from env.environment import DevOpsEnv
from env.models import VALID_ACTIONS
from graders.grader import compute_score
from tasks.task_config import TASK_CONFIGS

# ---------------------------------------------------------------------------
# Hackathon parameters
# ---------------------------------------------------------------------------
# ⚠️  IMPORTANT: whatever MODEL_ID you set here MUST be the same model you
# pass as --base-model in compare_inference.py.  Comparing a trained 1B
# against a base 3B is an unfair comparison — the 3B will always win on
# raw capability regardless of training quality.
#
# Recommended:  "unsloth/Llama-3.2-1B-Instruct"  (T4, ~30 min training)
#               "unsloth/Llama-3.2-3B-Instruct"  (T4, ~70 min training)
#               "unsloth/Meta-Llama-3.1-8B-Instruct"  (A100 / HF Endpoint)
# ---------------------------------------------------------------------------
MODEL_ID = "unsloth/Llama-3.2-1B-Instruct"   # change both here AND in compare_inference.py
MAX_SEQ_LENGTH = 1024
LORA_RANK = 32    # bumped from 16 → 32 for better expressivity at 1B scale
REWARD_LOG_PATH = "outputs_grpo/reward_log.csv"

system_prompt = (
    "You are an On-call SRE. Pick exactly one action from the valid actions.\n"
    "Valid actions: {}\n"
    "Do not explain your reasoning. Just output the action word."
).format(", ".join(VALID_ACTIONS))


# ---------------------------------------------------------------------------
# 1. Reward Functions (4 independent signals)
# ---------------------------------------------------------------------------

def format_reward_func(prompts, completions, **kwargs) -> List[float]:
    """Reward 1: Did the model output a single valid action string?

    Reduced from 1.0 → 0.3 so this baseline signal no longer dominates
    the task-specific rewards and the model has a real incentive to explore
    diagnostic / mitigation actions rather than any valid action.
    """
    rewards = []
    for completion in completions:
        text = completion[0]["content"].strip()
        rewards.append(0.3 if text in VALID_ACTIONS else -0.5)
    return rewards


def step_reward_func(prompts, completions, **kwargs) -> List[float]:
    """Reward 2: Does this action improve the incident state?

    Restores approximate mid-episode state by replaying recent_actions from the
    serialised observation in the user message before scoring the candidate action.
    This fixes the bug where mid-episode prompts were always evaluated from step 0.

    NaN safety: the entire env instantiation + step is wrapped in try/except.
    A bad task name, corrupted prompt, or any env error returns -0.5 instead
    of propagating an exception that would produce a NaN reward in the batch.
    The return value is also explicitly clamped to [-1.0, 1.0].
    """
    rewards = []
    for prompt, completion in zip(prompts, completions):
        action = completion[0]["content"].strip()

        if action not in VALID_ACTIONS:
            rewards.append(-0.5)
            continue

        try:
            task = "easy"
            recent_actions: List[str] = []
            try:
                state_dict = json.loads(prompt[-1]["content"])
                task = state_dict.get("task", "easy")
                if task not in TASK_CONFIGS:   # guard against unknown task names
                    task = "easy"
                recent_actions = state_dict.get("recent_actions", [])
            except Exception:
                pass

            env = DevOpsEnv(task=task)
            env.reset()
            # Replay the recent actions to reconstruct approximate mid-episode state
            for prev in recent_actions:
                if prev in VALID_ACTIONS:
                    env.step(prev)

            _, step_reward, _, _ = env.step(action)
            # Clamp to [-1, 1] — env already clamps to [-0.25, 1.0] but be explicit
            rewards.append(float(max(-1.0, min(1.0, step_reward))))

        except Exception:
            # Fallback: any env error → neutral negative reward, not NaN
            rewards.append(-0.5)

    return rewards


# Actions that give trivially safe small rewards — the model's previous local optimum
_FILLER_ACTIONS = {"acknowledge_incident", "post_status_update"}


def anti_cheat_reward_func(prompts, completions, **kwargs) -> List[float]:
    """Reward 3: Penalise reward hacking.

    BUG FIX (v2): The previous version read from `recent_actions` which the
    environment caps at the last 5 actions.  More critically, most GRPO
    training prompts come from *initial* episode states where recent_actions
    is [], so the repetition penalty (-0.8) never fired during training.

    This version reads `actions_taken` (the full episode history injected into
    mid-episode prompts by generate_prompts), falling back to recent_actions
    when actions_taken is not present.  The loop-detection branch fires an
    even stronger penalty when the last 3 actions are all the same token.
    """
    rewards = []
    for prompt, completion in zip(prompts, completions):
        action = completion[0]["content"].strip()

        all_actions_taken: List[str] = []
        try:
            state_dict = json.loads(prompt[-1]["content"])
            # Prefer full history; fall back to sliding window.
            all_actions_taken = state_dict.get(
                "actions_taken", state_dict.get("recent_actions", [])
            )
        except Exception:
            pass

        if action == "no_op":
            rewards.append(-0.5)          # strong no-op penalty
        elif action == "resolve_incident":
            rewards.append(-0.2)          # premature resolve penalty
        elif action in _FILLER_ACTIONS:
            rewards.append(-0.15)         # penalise spamming safe filler actions
        elif len(all_actions_taken) >= 3 and len(set(all_actions_taken[-3:])) == 1:
            # Last 3 actions are identical — loop detected, hard break
            rewards.append(-1.0)
        elif action in all_actions_taken:
            # Action already taken somewhere in the episode
            rewards.append(-0.8)
        else:
            rewards.append(0.4)           # reward any NEW genuine diagnostic/mitigation
    return rewards


def task_alignment_reward_func(prompts, completions, **kwargs) -> List[float]:
    """Reward 4: Strong bonus for picking the task-correct diagnostic or mitigation.

    Key fix: reward for required actions raised from 0.3 → 0.8, making the total
    reward for a correct diagnostic action ~1.4 vs ~-0.35 for a filler action.
    This gap is large enough for a 1B model to discover via GRPO exploration.
    Also ensures no reward is given for redundant actions already taken.
    """
    rewards = []
    for prompt, completion in zip(prompts, completions):
        action = completion[0]["content"].strip()

        if action not in VALID_ACTIONS:
            rewards.append(-0.3)
            continue

        task = "easy"
        recent_actions = []
        try:
            state_dict = json.loads(prompt[-1]["content"])
            task = state_dict.get("task", "easy")
            recent_actions = state_dict.get("recent_actions", [])
        except Exception:
            pass

        if action in recent_actions and action not in _FILLER_ACTIONS:
            # Remove any task alignment incentive to farm repeated actions
            rewards.append(-0.5)
            continue

        config = TASK_CONFIGS.get(task, TASK_CONFIGS["easy"])
        required_diag = config.get("required_diagnostics", [])
        required_mit  = config.get("required_mitigations", [])
        good_followups = config.get("good_followups", [])

        if action in required_diag or action in required_mit:
            rewards.append(0.8)    # massive reward for the right tool on the right task
        elif action in good_followups:
            rewards.append(0.3)    # good follow-up action
        elif action in _FILLER_ACTIONS:
            rewards.append(0.0)    # no alignment credit for filler
        elif action == "no_op":
            rewards.append(-0.5)
        else:
            rewards.append(-0.2)   # valid action but wrong task

    return rewards


def sequence_progress_reward_func(prompts, completions, **kwargs) -> List[float]:
    """Reward 5: Bonus when the agent has completed required diagnostics
    before attempting mitigations — reinforces the correct SRE workflow:
    investigate → mitigate → communicate → resolve.

    BUG FIX (v2): Now reads `actions_taken` (full history) in addition to
    `recent_actions` so that mid-episode prompts carry complete context.
    """
    rewards = []
    for prompt, completion in zip(prompts, completions):
        action = completion[0]["content"].strip()

        if action not in VALID_ACTIONS:
            rewards.append(0.0)
            continue

        task = "easy"
        all_actions_taken: List[str] = []
        try:
            state_dict = json.loads(prompt[-1]["content"])
            task = state_dict.get("task", "easy")
            all_actions_taken = state_dict.get(
                "actions_taken", state_dict.get("recent_actions", [])
            )
        except Exception:
            pass

        config = TASK_CONFIGS.get(task, TASK_CONFIGS["easy"])
        required_diag = set(config.get("required_diagnostics", []))
        required_mit  = set(config.get("required_mitigations", []))
        done_so_far   = set(all_actions_taken)

        diag_complete = required_diag.issubset(done_so_far)

        if action in required_mit and diag_complete:
            # Agent is attempting mitigation AFTER completing diagnosis — perfect sequence
            rewards.append(0.5)
        elif action in required_diag and not diag_complete:
            # Agent is still in the diagnostic phase — good
            rewards.append(0.2)
        elif action in required_mit and not diag_complete:
            # Attempting mitigation before diagnosis — penalise slightly
            rewards.append(-0.1)
        else:
            rewards.append(0.0)

    return rewards


def diversity_reward_func(prompts, completions, **kwargs) -> List[float]:
    """Reward 6 (NEW): Penalise GRPO group-level mode collapse.

    GRPO samples `num_generations` completions per prompt and learns from
    relative advantage.  When all completions are identical the reward_std
    is 0 and there is NO gradient signal — the model stops learning entirely.

    This function detects that collapse at the batch level and applies a
    uniform negative signal so the KL term pushes the policy back toward
    exploration.  It returns 0.0 when the group shows healthy diversity.
    """
    actions = [c[0]["content"].strip() for c in completions]
    valid_actions = [a for a in actions if a in VALID_ACTIONS]
    unique_count = len(set(valid_actions)) if valid_actions else 0

    if unique_count <= 1:
        # Complete collapse — penalise uniformly to force exploration
        return [-0.25] * len(completions)
    else:
        return [0.0] * len(completions)


# ---------------------------------------------------------------------------
# 2. Reward Curve Logger Callback
# ---------------------------------------------------------------------------

try:
    from transformers import TrainerCallback as _TrainerCallbackBase
except ImportError:  # allow import on CPU-only machines without transformers
    _TrainerCallbackBase = object  # type: ignore[misc,assignment]


class RewardLoggerCallback(_TrainerCallbackBase):  # type: ignore[valid-type]
    """Saves per-step reward metrics to a CSV for demo reward curves.

    Handles both TRL naming conventions for per-function reward keys:
      - Old TRL (< ~0.15):  reward_format_reward_func   (underscore prefix)
      - New TRL (>= ~0.15): rewards/format_reward_func  (slash prefix)

    All keys are normalised to the underscore form before writing so the
    CSV column names are stable regardless of which TRL version is installed.
    Only rows where 'reward' is present are written — non-training log
    events (LR warmup, checkpoint saves) are skipped.
    """

    # Canonical CSV column names.
    CSV_COLUMNS = [
        "step",
        "reward",
        "reward_format_reward_func",
        "reward_step_reward_func",
        "reward_anti_cheat_reward_func",
        "reward_task_alignment_reward_func",
        "reward_sequence_progress_reward_func",
    ]

    def __init__(self, log_path: str = REWARD_LOG_PATH):
        self.log_path = log_path
        self._header_written = False
        os.makedirs(os.path.dirname(log_path), exist_ok=True)

    @staticmethod
    def _normalize(key: str) -> str:
        """Convert TRL's slash-format to underscore-format.

        rewards/format_reward_func  ->  reward_format_reward_func
        reward_format_reward_func   ->  reward_format_reward_func (unchanged)
        """
        if key.startswith("rewards/"):
            return "reward_" + key[len("rewards/"):]
        return key

    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs is None or "reward" not in logs:
            return

        # Capture every reward-related key TRL emits and normalize it.
        row: dict = {"step": state.global_step}
        for raw_key, value in logs.items():
            if raw_key == "reward" or raw_key.startswith("reward_") or raw_key.startswith("rewards/"):
                col = self._normalize(raw_key)
                row[col] = value

        # Fill any canonical column that TRL didn't emit this step.
        for col in self.CSV_COLUMNS:
            row.setdefault(col, "")

        write_header = not self._header_written and not os.path.exists(self.log_path)
        with open(self.log_path, "a", newline="") as f:
            writer = csv.DictWriter(
                f, fieldnames=self.CSV_COLUMNS, extrasaction="ignore"
            )
            if write_header:
                writer.writeheader()
                self._header_written = True
            writer.writerow(row)


# ---------------------------------------------------------------------------
# 3. Curriculum Dataset Generation
# ---------------------------------------------------------------------------

def _make_prompt(state: dict) -> dict:
    return {
        "prompt": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": json.dumps(state)},
        ]
    }


def generate_sft_dataset(seed: int = 42) -> Dataset:
    """Generate a supervised fine-tuning dataset from optimal action sequences.

    For each of the 6 tasks we replay the full optimal trajectory and record
    every (observation, optimal_next_action) pair as a prompt/completion.
    This gives the model a strong starting policy before GRPO kicks in.
    """
    random.seed(seed)
    data = []

    for task_name, config in TASK_CONFIGS.items():
        optimal_actions = config.get("optimal_actions", [])
        if not optimal_actions:
            continue

        env = DevOpsEnv(task=task_name)
        state = env.reset()

        for action in optimal_actions:
            prompt = [
                {"role": "system", "content": system_prompt},
                {"role": "user",   "content": json.dumps(state)},
            ]
            # SFT target: the optimal action at this point in the episode
            data.append({
                "prompt": prompt,
                "completion": [{"role": "assistant", "content": action}],
            })
            state, _, done, _ = env.step(action)
            if done:
                break

    random.shuffle(data)
    print(f"  SFT dataset: {len(data)} (state, optimal_action) pairs across all 6 tasks.")
    return Dataset.from_list(data)


def _make_prompt_with_history(state: dict, actions_taken: List[str]) -> dict:
    """Like _make_prompt but injects the full actions_taken list so that
    anti_cheat_reward_func and sequence_progress_reward_func can correctly
    penalise repeated actions in mid-episode prompts.

    BUG FIX: Previously mid-episode prompts only carried `recent_actions`
    (the last-5 sliding window), so the anti-cheat penalty never fired for
    actions taken more than 5 steps ago, and the repetition check was always
    empty for initial-state prompts.
    """
    state_with_history = {**state, "actions_taken": actions_taken}
    return {
        "prompt": [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": json.dumps(state_with_history)},
        ]
    }


def generate_prompts(
    per_task_n: int = 15,
    mid_episode_n: int = 15,
    seed: int = 42,
    easy_n: int = None,
    medium_n: int = None,
    hard_n: int = None,
) -> Dataset:
    """Generate a curriculum-mixed prompt dataset for GRPO.

    Distribution:
      - per_task_n prompts per task initial state (all 6 tasks)
        OR easy_n/medium_n/hard_n for per-task control (easy/medium/hard only)
      - mid_episode_n : states captured after 1-2 random valid steps

    BUG FIX: Mid-episode prompts now include `actions_taken` (the full episode
    history) via _make_prompt_with_history so that reward functions can penalise
    repeated actions correctly.  Initial-state prompts have an empty
    actions_taken by definition and are unchanged.
    """
    random.seed(seed)
    data = []

    # Support per-task counts (easy_n / medium_n / hard_n) or fall back to per_task_n
    per_task_counts = {}
    if easy_n is not None or medium_n is not None or hard_n is not None:
        per_task_counts = {
            "easy": easy_n if easy_n is not None else per_task_n,
            "medium": medium_n if medium_n is not None else per_task_n,
            "hard": hard_n if hard_n is not None else per_task_n,
        }
        all_tasks = list(per_task_counts.keys())
    else:
        all_tasks = list(TASK_CONFIGS.keys())
        per_task_counts = {task: per_task_n for task in all_tasks}

    # Initial states — actions_taken is empty at episode start, so no history needed
    for task in all_tasks:
        env = DevOpsEnv(task=task)
        state = env.reset()
        for _ in range(per_task_counts[task]):
            data.append(_make_prompt(state))

    # Mid-episode states — inject full action history for correct anti-cheat scoring
    for _ in range(mid_episode_n):
        task = random.choice(all_tasks)
        n_warm = random.randint(1, 3)
        env = DevOpsEnv(task=task)
        state = env.reset()
        for _ in range(n_warm):
            action = random.choice(VALID_ACTIONS)
            state, _, done, _ = env.step(action)
            if done:
                state = env.reset()
                break
        # Pass the full actions_taken so reward functions see real episode history
        data.append(_make_prompt_with_history(state, env._state["actions_taken"]))

    random.shuffle(data)
    print(f"  GRPO dataset: {len(data)} prompts across {len(all_tasks)} tasks + mid-episode states.")
    return Dataset.from_list(data)


# ---------------------------------------------------------------------------
# 4. Main Training Loop
# ---------------------------------------------------------------------------

def main():
    # GPU-only imports — kept here so the module can be imported without a GPU
    import torch  # noqa: F401
    from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
    from unsloth import FastLanguageModel, PatchDPOTrainer
    from trl import GRPOTrainer, GRPOConfig, SFTTrainer, SFTConfig

    print("[1] Loading Unsloth Model...")
    PatchDPOTrainer()  # Required optimisations

    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_ID,
        max_seq_length=MAX_SEQ_LENGTH,
        dtype=None,
        load_in_4bit=True,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_RANK,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        lora_alpha=16,
        use_gradient_checkpointing="unsloth",
    )

    # ── Phase 1: SFT Warm-start ────────────────────────────────────────────
    # Run 1 epoch of supervised fine-tuning on optimal action sequences.
    # This gives GRPO a non-random starting policy so rewards are non-zero
    # from the very first update step.
    print("[2] Phase 1 — SFT Warm-start on optimal trajectories...")
    sft_dataset = generate_sft_dataset()

    def format_sft_sample(example):
        """Convert prompt+completion into a single flat text for SFT."""
        messages = example["prompt"] + example["completion"]
        return {"text": tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=False
        )}

    sft_dataset = sft_dataset.map(format_sft_sample)

    sft_args = SFTConfig(
        output_dir="outputs_sft",
        num_train_epochs=1,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        logging_steps=1,
        dataset_text_field="text",
        max_seq_length=MAX_SEQ_LENGTH,
        fp16=not torch.cuda.is_bf16_supported(),
        bf16=torch.cuda.is_bf16_supported(),
    )

    sft_trainer = SFTTrainer(
        model=model,
        tokenizer=tokenizer,
        args=sft_args,
        train_dataset=sft_dataset,
    )

    print("  Running SFT warm-start epoch...")
    sft_trainer.train()
    print("  SFT warm-start complete.")

    # ── Phase 2: GRPO RL Training ──────────────────────────────────────────
    print("[3] Phase 2 — Generating GRPO Curriculum (all 6 tasks)...")
    grpo_dataset = generate_prompts()

    print("[4] Setting up GRPO Config (max_steps=300)...")
    training_args = GRPOConfig(
        output_dir="outputs_grpo",
        learning_rate=5e-6,              # halved from 1e-5 — reduces KL explosion risk
        lr_scheduler_type="cosine",
        warmup_steps=20,                 # NaN fix: let optimizer calibrate before full LR kicks in
        max_steps=300,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        logging_steps=1,
        # NaN prevention — the three settings below address the KL spike → NaN chain
        max_grad_norm=0.1,               # aggressive clip (default=1.0); stops gradient overflow when KL spikes
        # GRPO-specific
        num_generations=8,               # more group samples → reward_std stays non-zero even during partial collapse
        max_prompt_length=512,
        max_completion_length=32,
        temperature=0.9,                 # adds generation stochasticity so completions are not all identical
    )

    trainer = GRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=[
            format_reward_func,
            step_reward_func,
            anti_cheat_reward_func,
            task_alignment_reward_func,
            sequence_progress_reward_func,
            diversity_reward_func,       # NEW: breaks GRPO variance collapse
        ],
        args=training_args,
        train_dataset=grpo_dataset,
        callbacks=[RewardLoggerCallback()],
    )

    print("[5] Beginning GRPO RL Training Loop...")
    trainer.train()

    print("[6] Saving merged model...")
    model.save_pretrained_merged("trained_sre_agent", tokenizer, save_method="merged_16bit")
    print(f"Training Complete! Reward log saved to {REWARD_LOG_PATH}")


if __name__ == "__main__":
    main()