training/train_grpo.py

"""
GRPO Training Loop — Fine-tunes the DevOps agent using Group Relative Policy Optimization.

Uses TRL's GRPOTrainer with Unsloth for efficient LoRA fine-tuning.
Integrates curriculum scheduler (rolling windows), replay buffer,
anti-reward-hacking checks, and proper LoRA weight saving.

Per the hackathon guide:
  - Build the environment FIRST. Do not touch the trainer until reset/step/rewards
    are locally verified and stable.
  - Actively guard against reward hacking.
  - Save LoRA/QLoRA weights correctly. Do NOT upcast 4-bit to 16-bit before merging.
  - Inspect actual generations during training — do not rely only on mean reward.
"""

from __future__ import annotations

import json
import os
import inspect
import time
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

from agent.baseline_agent import BaselineAgent
from agent.prompts import format_chat_messages, format_prompt
from devops_env.env import DevOpsEnv
from replay.buffer import ReplayBuffer
from scenarios.registry import ScenarioRegistry
from training.curriculum import CurriculumScheduler


class AntiHackingMonitor:
    """Monitors for reward hacking patterns during training.

    Checks:
      1. Overall reward rising while success rate stays flat → likely hacking
      2. Repeated commands across episodes (cached/memorized outputs)
      3. Dangerous command reward firing more than once per run
      4. Success column not moving despite total reward increase

    Usage:
        monitor = AntiHackingMonitor()
        monitor.record_episode(episode_data)
        alerts = monitor.check()
    """

    def __init__(self, alert_threshold: int = 50) -> None:
        """Initialize the anti-hacking monitor.

        Args:
            alert_threshold: Check for hacking every N episodes.
        """
        self.alert_threshold = alert_threshold
        self._reward_history: List[float] = []
        self._success_history: List[bool] = []
        self._dangerous_count: int = 0
        self._generation_samples: List[Dict] = []
        self._command_frequency: Dict[str, int] = defaultdict(int)

    def record_episode(self, episode_data: Dict) -> None:
        """Record an episode's data for monitoring.

        Args:
            episode_data: Dict with total_reward, solved, steps, etc.
        """
        self._reward_history.append(episode_data.get("total_reward", 0.0))
        self._success_history.append(episode_data.get("solved", False))

        for step in episode_data.get("steps", []):
            action = step.get("action", "")
            self._command_frequency[action] += 1

            breakdown = step.get("reward_breakdown", {})
            if "dangerous_command" in breakdown:
                self._dangerous_count += 1

        # Sample generation for inspection
        if len(self._reward_history) % self.alert_threshold == 0:
            self._generation_samples.append({
                "episode": len(self._reward_history),
                "scenario": episode_data.get("scenario_id", ""),
                "commands": [s.get("action", "") for s in episode_data.get("steps", [])],
                "solved": episode_data.get("solved", False),
                "reward": episode_data.get("total_reward", 0.0),
            })

    def check(self) -> List[str]:
        """Run all anti-hacking checks.

        Returns:
            List of alert messages. Empty list = no issues detected.
        """
        alerts = []

        # Check 1: Reward rising but success flat
        if len(self._reward_history) >= 100:
            recent_50_reward = sum(self._reward_history[-50:]) / 50
            older_50_reward = sum(self._reward_history[-100:-50]) / 50
            recent_50_success = sum(self._success_history[-50:]) / 50
            older_50_success = sum(self._success_history[-100:-50]) / 50

            reward_increase = recent_50_reward - older_50_reward
            success_change = recent_50_success - older_50_success

            if reward_increase > 2.0 and success_change < 0.05:
                alerts.append(
                    f"⚠ REWARD HACKING SUSPECTED: Mean reward increased by "
                    f"{reward_increase:.1f} but success rate only changed by "
                    f"{success_change:.1%}. Check for environment exploits."
                )

        # Check 2: Dangerous commands firing too often
        if self._dangerous_count > 3:
            alerts.append(
                f"⚠ DANGEROUS COMMANDS: {self._dangerous_count} dangerous command "
                f"penalties detected. Agent may be probing blocklist boundaries."
            )

        # Check 3: Suspiciously repeated commands across episodes
        top_commands = sorted(
            self._command_frequency.items(), key=lambda x: x[1], reverse=True
        )[:5]
        total_commands = sum(self._command_frequency.values())
        if total_commands > 50 and top_commands:
            top_freq = top_commands[0][1] / total_commands
            if top_freq > 0.5:
                alerts.append(
                    f"⚠ COMMAND REPETITION: '{top_commands[0][0]}' used in "
                    f"{top_freq:.0%} of all commands. Possible memorization."
                )

        return alerts

    def get_generation_samples(self) -> List[Dict]:
        """Get sampled generations for manual inspection.

        Returns:
            List of generation sample dicts.
        """
        return self._generation_samples

    def print_sample_report(self) -> None:
        """Print the latest generation samples to console for inspection."""
        if not self._generation_samples:
            return

        print("\n" + "=" * 60)
        print("  GENERATION INSPECTION SAMPLES")
        print("=" * 60)

        for sample in self._generation_samples[-3:]:
            solved_str = "✓ SOLVED" if sample["solved"] else "✗ FAILED"
            print(f"\n  Episode {sample['episode']} | {sample['scenario']} | {solved_str}")
            print(f"  Reward: {sample['reward']:+.1f}")
            for i, cmd in enumerate(sample["commands"], 1):
                print(f"    Step {i}: {cmd}")

        print("=" * 60 + "\n")


class GRPODevOpsTrainer:
    """GRPO training loop for the DevOps RL agent.

    Runs rollout episodes, collects (prompt, completion, reward) tuples,
    and trains the model using TRL's GRPO approach with grouped completions.

    Includes:
      - Curriculum learning with rolling 50-episode windows
      - Anti-reward-hacking monitoring
      - Generation sample inspection every 50 steps
      - Proper LoRA weight saving (no 4-bit → 16-bit upcast)

    Usage:
        trainer = GRPODevOpsTrainer(model_name="unsloth/llama-3.2-3b-instruct")
        trainer.train(num_episodes=500)
    """

    def __init__(
        self,
        model_name: str = "unsloth/llama-3.2-3b-instruct",
        output_dir: str = "./checkpoints",
        db_url: str = "sqlite:///training_replay.db",
        num_generations: int = 4,
        max_new_tokens: int = 64,
        temperature: float = 0.8,
        learning_rate: float = 5e-5,
        batch_size: int = 4,
        gradient_accumulation_steps: int = 4,
        max_steps: int = 1000,
        save_steps: int = 100,
        logging_steps: int = 10,
    ) -> None:
        """Initialize the GRPO trainer.

        Args:
            model_name: HuggingFace model ID for the base model.
            output_dir: Directory for checkpoints.
            db_url: SQLAlchemy URL for the replay buffer.
            num_generations: Number of completions per prompt (GRPO needs groups).
            max_new_tokens: Max tokens per generation.
            temperature: Sampling temperature during rollouts.
            learning_rate: Learning rate for fine-tuning.
            batch_size: Per-device training batch size.
            gradient_accumulation_steps: Gradient accumulation factor.
            max_steps: Total training steps.
            save_steps: Save checkpoint every N steps.
            logging_steps: Log metrics every N steps.
        """
        self.model_name = model_name
        self.output_dir = output_dir
        self.num_generations = num_generations
        self.max_new_tokens = max_new_tokens
        self.temperature = temperature
        self.learning_rate = learning_rate
        self.batch_size = batch_size
        self.gradient_accumulation_steps = gradient_accumulation_steps
        self.max_steps = max_steps
        self.save_steps = save_steps
        self.logging_steps = logging_steps

        # Components
        self.replay_buffer = ReplayBuffer(db_url)
        self.curriculum = CurriculumScheduler(unlock_threshold=0.8, window_size=50)
        self.anti_hacking = AntiHackingMonitor(alert_threshold=50)
        self.registry = ScenarioRegistry()
        self.registry.register_defaults()

        # Model state
        self._model = None
        self._tokenizer = None
        self._trainer = None

        # Reward breakdown tracking (log each column separately)
        self._reward_column_totals: Dict[str, List[float]] = defaultdict(list)

    def _setup_model(self) -> None:
        """Load and prepare the model with LoRA for GRPO training.

        WARNING: Uses Unsloth's native 4-bit loading. Do NOT upcast
        4-bit to 16-bit before merging — this damages quality.
        """
        try:
            from unsloth import FastLanguageModel

            model, tokenizer = FastLanguageModel.from_pretrained(
                model_name=self.model_name,
                max_seq_length=2048,
                load_in_4bit=True,
                dtype=None,
            )

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

            self._model = model
            self._tokenizer = tokenizer
            print(f"[Trainer] ✓ Model loaded: {self.model_name}")
            print(f"[Trainer]   LoRA rank=16, 4-bit quantization enabled")

        except ImportError:
            print("[Trainer] ⚠ Unsloth not available. Trying transformers fallback...")
            try:
                from transformers import AutoModelForCausalLM, AutoTokenizer
                self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
                self._model = AutoModelForCausalLM.from_pretrained(
                    self.model_name, device_map="auto",
                )
                print(f"[Trainer] ✓ Model loaded via transformers: {self.model_name}")
            except Exception as e:
                print(f"[Trainer] ✗ Model load failed: {e}")
                print(f"[Trainer]   Will use baseline (rule-based) agent for rollouts")

    def _setup_grpo_trainer(self) -> None:
        """Configure the TRL GRPO trainer."""
        from trl import GRPOTrainer, GRPOConfig

        config_kwargs = {
            "output_dir": self.output_dir,
            "num_generations": self.num_generations,
            "max_new_tokens": self.max_new_tokens,
            "learning_rate": self.learning_rate,
            "per_device_train_batch_size": self.batch_size,
            "gradient_accumulation_steps": self.gradient_accumulation_steps,
            "max_steps": self.max_steps,
            "save_steps": self.save_steps,
            "logging_steps": self.logging_steps,
            "report_to": "none",
        }

        # TRL API changed in >=0.12; pass generation temperature only when supported.
        params = inspect.signature(GRPOConfig.__init__).parameters
        if "temperature" in params:
            config_kwargs["temperature"] = self.temperature
        elif "generation_kwargs" in params:
            config_kwargs["generation_kwargs"] = {"temperature": self.temperature}

        config = GRPOConfig(**config_kwargs)

        self._trainer = GRPOTrainer(
            model=self._model,
            processing_class=self._tokenizer,
            config=config,
            reward_funcs=self._reward_function,
        )

    def _reward_function(self, completions: List[str], **kwargs) -> List[float]:
        """Compute rewards for a batch of GRPO completions.

        Args:
            completions: List of generated shell commands.

        Returns:
            List of reward values.
        """
        rewards = []

        level = kwargs.get("level")
        if isinstance(level, list):
            level = level[0] if level else None
        if level is None:
            level = self.curriculum.sample_level()

        scenario_id = kwargs.get("scenario_id")
        if isinstance(scenario_id, list):
            scenario_id = scenario_id[0] if scenario_id else None
        if not scenario_id:
            scenario_id = self.registry.get_random(level=level).id

        for completion in completions:
            command = completion.strip()
            env = None
            try:
                # All completions in a group must be evaluated on the same scenario.
                env = DevOpsEnv(
                    scenario_registry=self.registry,
                    max_steps=1,
                    target_level=level,
                    target_scenario=scenario_id,
                )
                env.reset(options={"scenario_id": scenario_id})
                _, reward, _, _, _ = env.step(command)
                rewards.append(reward)
            except Exception:
                rewards.append(-1.0)
            finally:
                if env is not None:
                    env.close()
        return rewards

    def run_rollout_episode(self, level: int | None = None) -> Dict:
        """Run a single rollout episode using the current agent.

        Uses the LLM agent if loaded, otherwise falls back to the
        rule-based baseline agent.

        Args:
            level: Difficulty level to use. If None, curriculum decides.

        Returns:
            Episode summary dict.
        """
        # Use baseline agent if model not loaded
        if self._model is not None:
            from agent.devops_agent import DevOpsAgent
            agent = DevOpsAgent(
                model_name=self.model_name,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
                model=self._model,
                tokenizer=self._tokenizer,
                auto_load=False,
            )
        else:
            agent = BaselineAgent()

        selected_level = level if level is not None else self.curriculum.sample_level()
        env = DevOpsEnv(
            scenario_registry=self.registry,
            target_level=selected_level,
        )

        obs, info = env.reset()
        total_reward = 0.0
        steps = []
        done = False

        while not done:
            action = agent.act(obs)
            obs, reward, terminated, truncated, step_info = env.step(action)
            total_reward += reward

            step_data = {
                "step": step_info.get("step_count", len(steps) + 1),
                "action": action,
                "reward": reward,
                "reward_breakdown": step_info.get("reward_breakdown", {}),
                "error_type": obs.get("error_type", "unknown"),
                "observation": {
                    "error_log": obs.get("error_log", "")[:500],
                    "command_history": obs.get("command_history", []),
                    "step_count": obs.get("step_count", 0),
                },
                "result": step_info.get("execution_result", {}),
            }
            steps.append(step_data)

            # Track individual reward columns
            for col, val in step_info.get("reward_breakdown", {}).items():
                self._reward_column_totals[col].append(val)

            done = terminated or truncated

        summary = env.get_episode_summary()
        env.close()

        # Store in replay buffer
        episode_id = self.replay_buffer.store_episode(
            scenario_id=summary["scenario_id"],
            level=summary["level"],
            steps=steps,
            total_reward=total_reward,
            solved=summary["solved"],
        )

        summary["episode_id"] = episode_id
        summary["steps"] = steps
        return summary

    def train(self, num_episodes: int = 500, use_grpo: bool = True) -> Dict:
        """Run the full training loop.

        Args:
            num_episodes: Total number of rollout episodes.
            use_grpo: Whether to use GRPO training (requires GPU + Unsloth).

        Returns:
            Training summary with metrics.
        """
        print(f"\n{'='*60}")
        print(f"  GRPO Training — DevOps RL Agent")
        print(f"{'='*60}")
        print(f"  Model: {self.model_name}")
        print(f"  Episodes: {num_episodes}")
        print(f"  Curriculum: {self.curriculum.get_status()}")
        print(f"{'='*60}\n")

        if use_grpo and self._model is None:
            try:
                self._setup_model()
                self._setup_grpo_trainer()
            except Exception as e:
                print(f"[Trainer] ⚠ GRPO setup failed: {e}")
                print(f"[Trainer]   Running rollouts with baseline agent only.")
                use_grpo = False

        metrics_history = []

        for episode_num in range(num_episodes):
            level = self.curriculum.sample_level()
            summary = self.run_rollout_episode(level=level)

            # Record in curriculum (rolling window)
            self.curriculum.record_episode(
                level=summary["level"],
                solved=summary["solved"],
            )

            # Record for anti-hacking monitoring
            self.anti_hacking.record_episode(summary)

            # Periodic logging
            if (episode_num + 1) % self.logging_steps == 0:
                metrics = self._compute_metrics(episode_num + 1, summary)
                metrics_history.append(metrics)
                self._print_progress(metrics)

            # Inspect actual generations every 50 episodes
            if (episode_num + 1) % 50 == 0:
                self.anti_hacking.print_sample_report()

                # Run anti-hacking checks
                alerts = self.anti_hacking.check()
                for alert in alerts:
                    print(f"\n  {alert}\n")

                # Log reward column breakdown
                self._print_reward_column_breakdown()

            # Save checkpoint periodically
            if use_grpo and self._model and (episode_num + 1) % self.save_steps == 0:
                self._save_checkpoint(episode_num + 1)

                # Test post-training inference immediately after save
                self._verify_checkpoint(episode_num + 1)

        final_stats = self.replay_buffer.get_stats()
        print(f"\n{'='*60}")
        print(f"  TRAINING COMPLETE")
        print(f"{'='*60}")
        print(f"  Total episodes: {num_episodes}")
        print(f"  Curriculum status: {self.curriculum.get_status()}")
        print(json.dumps(final_stats, indent=2))

        return {
            "total_episodes": num_episodes,
            "final_stats": final_stats,
            "metrics_history": metrics_history,
            "anti_hacking_alerts": self.anti_hacking.check(),
        }

    def _compute_metrics(self, episode_num: int, latest_summary: Dict) -> Dict:
        """Compute training metrics at a logging step."""
        status = self.curriculum.get_status()
        return {
            "episode": episode_num,
            "scenario": latest_summary.get("scenario_id", ""),
            "solved": latest_summary.get("solved", False),
            "reward": latest_summary.get("total_reward", 0.0),
            "steps": latest_summary.get("total_steps", 0),
            "curriculum": status,
            "l1_solve_rate": status[1]["window_solve_rate"],
            "l2_solve_rate": status[2]["window_solve_rate"],
            "l3_solve_rate": status[3]["window_solve_rate"],
        }

    def _print_progress(self, metrics: Dict) -> None:
        """Print training progress to console."""
        ep = metrics["episode"]
        solved = "✓" if metrics["solved"] else "✗"
        reward = metrics["reward"]
        scenario = metrics["scenario"]

        level_info = []
        for lvl in [1, 2, 3]:
            status = metrics["curriculum"][lvl]
            if status["unlocked"]:
                rate = status["window_solve_rate"]
                eps = status["total_episodes"]
                level_info.append(f"L{lvl}:{rate:.0%}({eps})")

        print(f"  [{ep:>4d}] {scenario:<22s} {solved} r={reward:>6.1f} | {'  '.join(level_info)}")

    def _print_reward_column_breakdown(self) -> None:
        """Print per-column reward averages for hacking detection."""
        if not self._reward_column_totals:
            return

        print("\n  Reward Column Breakdown (last window):")
        for col, values in sorted(self._reward_column_totals.items()):
            recent = values[-50:] if len(values) >= 50 else values
            avg = sum(recent) / len(recent) if recent else 0
            direction = "↑" if avg > 0 else "↓" if avg < 0 else "—"
            print(f"    {col:<20s}: {avg:>+6.2f} {direction}")
        print()

    def _save_checkpoint(self, step: int) -> None:
        """Save LoRA adapter weights correctly.

        WARNING: Do NOT upcast 4-bit to 16-bit and naively merge.
        Save adapters directly using save_pretrained.
        """
        ckpt_dir = os.path.join(self.output_dir, f"checkpoint-{step}")
        os.makedirs(ckpt_dir, exist_ok=True)
        try:
            if self._model is not None:
                # Save LoRA adapters directly — NOT merged with base model
                self._model.save_pretrained(ckpt_dir)
                if self._tokenizer is not None:
                    self._tokenizer.save_pretrained(ckpt_dir)
                print(f"[Trainer] ✓ Checkpoint saved: {ckpt_dir}")
                print(f"[Trainer]   Saved as LoRA adapters (not merged)")
        except Exception as e:
            print(f"[Trainer] ✗ Checkpoint save failed: {e}")

    def _verify_checkpoint(self, step: int) -> None:
        """Test post-training inference immediately after checkpoint save.

        Per the hackathon guide: "Test post-training inference immediately
        after export, not at the end."
        """
        ckpt_dir = os.path.join(self.output_dir, f"checkpoint-{step}")
        try:
            from agent.devops_agent import DevOpsAgent
            if not os.path.isdir(ckpt_dir):
                print(f"[Trainer] ✗ Post-save verification failed: missing {ckpt_dir}")
                return

            adapter_files = [
                "adapter_config.json",
                "adapter_model.safetensors",
                "adapter_model.bin",
            ]
            if not any(os.path.exists(os.path.join(ckpt_dir, f)) for f in adapter_files):
                print(f"[Trainer] ✗ Post-save verification failed: adapter files missing in {ckpt_dir}")
                return

            test_agent = DevOpsAgent(
                model_name=self.model_name,
                max_new_tokens=self.max_new_tokens,
                temperature=self.temperature,
            )

            if test_agent.model_name == "rule-based" or not getattr(test_agent, "_is_loaded", False):
                print("[Trainer] ✗ Post-save verification failed: could not load base model")
                return

            test_agent.load_checkpoint(ckpt_dir)

            test_obs = {
                "error_log": "ModuleNotFoundError: No module named 'flask'",
                "command_history": [],
                "step_count": 0,
                "scenario_id": "missing_flask",
                "error_type": "missing_package",
            }
            result = test_agent.act(test_obs)
            if result:
                print(f"[Trainer] ✓ Post-save inference verified: '{result}'")
            else:
                print(f"[Trainer] ⚠ Post-save inference returned empty result")
        except Exception as e:
            print(f"[Trainer] ✗ Post-save inference check failed: {e}")