train.py

"""GRPO training pipeline for Enterprise Incident Command Center.

Runs on Google Colab free tier (T4 GPU) when optional training dependencies
are installed. Supports a deterministic dry-run mode for local verification.

Usage:
    python train.py --iterations 1 --episodes 1 --k 2 --dry-run
    python train.py --iterations 20 --episodes 30 --k 4
"""

from __future__ import annotations

import ast
import argparse
import asyncio
import inspect
import json
import random
import re
import subprocess
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

from env.environment import CustomerSupportEnv
from evaluate import (
    PolicyState,
    behavior_diffs,
    evaluate_policy,
    plot_reports,
)
from evaluate import choose_policy_action as eval_choose_policy_action
from models.observation import Observation


KNOWN_ACTION_TYPES: tuple[str, ...] = (
    "classify",
    "route",
    "respond",
    "escalate",
    "request_info",
    "resolve",
    "check_monitoring",
    "probe_service",
    "fetch_logs",
    "fetch_user_data",
    "check_billing",
    "query_kb",
    "check_policy",
    "query_incident_history",
    "follow_runbook_step",
    "apply_fix",
    "verify_fix",
    "rollback_fix",
    "notify_stakeholders",
    "write_postmortem",
    "update_kb",
)

_JSON_OBJECT_RE = re.compile(r"\{(?:[^{}]|(?:\{[^{}]*\}))*\}", re.DOTALL)


def _extract_first_json_object(text: str) -> dict[str, object] | None:
    """Return the first valid JSON object found in a completion.

    Tolerates chat prose, code fences, and extra whitespace so noisy model
    outputs still produce a usable reward signal instead of a degenerate
    constant penalty.
    """
    if not text:
        return None
    stripped = text.strip()
    try:
        payload = json.loads(stripped)
        if isinstance(payload, dict):
            return payload
    except json.JSONDecodeError:
        pass
    for match in _JSON_OBJECT_RE.finditer(text):
        try:
            payload = json.loads(match.group(0))
        except json.JSONDecodeError:
            continue
        if isinstance(payload, dict):
            return payload
    return None


def _extract_json_object_matches(text: str) -> list[re.Match[str]]:
    """Return regex matches for JSON-like objects found in text."""
    return list(_JSON_OBJECT_RE.finditer(text or ""))


@dataclass(slots=True)
class TrajectoryRow:
    """One prompt/completion/reward row used for GRPO datasets."""

    prompt: str
    completion: str
    reward: float
    iteration: int
    episode: int
    step: int
    difficulty: str


def curriculum_difficulty(iteration: int, episode_index: int, episodes: int) -> str:
    """Deterministic curriculum schedule from phase 7 specification."""
    if iteration <= 8:
        return "easy"
    if iteration <= 14:
        return "medium"
    if iteration <= 18:
        return "hard"

    # Phase D mixed schedule: easy 20% / medium 30% / hard 40% / nightmare 10%.
    slot = episode_index % max(1, episodes)
    ratio = slot / max(1, episodes)
    if ratio < 0.20:
        return "easy"
    if ratio < 0.50:
        return "medium"
    if ratio < 0.90:
        return "hard"
    return "nightmare"


_FORMAT_INSTRUCTION = (
    "You are an incident response agent. "
    "Respond with exactly ONE compact JSON object and nothing else. "
    'Example: {"action_type":"check_monitoring","service_name":null}. '
    "Pick action_type strictly from available_actions below. "
    "Include the fields that action requires."
)


def build_prompt(obs: Observation) -> str:
    """Convert incident observation into a deterministic training prompt.

    The prompt is structured so a downstream reward function can parse it
    back (phase, available_actions) and so the model sees an explicit
    JSON-only instruction with a concrete example.
    """
    parts = [
        _FORMAT_INSTRUCTION,
        f"incident={obs.incident_id}",
        f"title={obs.incident_title}",
        f"phase={obs.incident_phase}",
        f"step={obs.current_step}/{obs.max_steps}",
        f"available_actions={obs.available_actions}",
    ]
    if obs.system_status:
        parts.append(f"system_status={json.dumps(obs.system_status, sort_keys=True)}")
    if obs.tool_results:
        parts.append(f"tool_results={json.dumps(obs.tool_results, sort_keys=True)}")
    if obs.known_facts:
        parts.append(f"known_facts={json.dumps(obs.known_facts, sort_keys=True)}")
    parts.append("Respond with ONE JSON action object only:")
    return "\n".join(parts)


def choose_training_action(
    obs: Observation,
    state: PolicyState,
    quality_ratio: float,
) -> dict[str, object]:
    """Interpolation policy used to collect trajectories during curriculum."""
    if quality_ratio < 0.5:
        return eval_choose_policy_action(obs, state, "baseline")
    return eval_choose_policy_action(obs, state, "trained")


async def collect_trajectories(
    *,
    iterations: int,
    episodes: int,
) -> tuple[list[TrajectoryRow], list[float]]:
    """Collect deterministic trajectories across curriculum iterations."""
    env = CustomerSupportEnv()
    rows: list[TrajectoryRow] = []
    reward_history: list[float] = []

    try:
        for iteration in range(1, iterations + 1):
            quality_ratio = iteration / max(1, iterations)
            cumulative = 0.0
            reward_count = 0

            for episode_idx in range(episodes):
                difficulty = curriculum_difficulty(iteration, episode_idx, episodes)
                reset = await env.reset(
                    seed=episode_idx,
                    difficulty=difficulty,
                    mode="incident",
                )
                obs = reset.observation
                policy_state = PolicyState()

                for step_idx in range(obs.max_steps):
                    prompt = build_prompt(obs)
                    action = choose_training_action(obs, policy_state, quality_ratio)
                    completion = json.dumps(action, separators=(",", ":"))
                    result = await env.step(action)

                    rows.append(
                        TrajectoryRow(
                            prompt=prompt,
                            completion=completion,
                            reward=result.reward,
                            iteration=iteration,
                            episode=episode_idx,
                            step=step_idx,
                            difficulty=difficulty,
                        )
                    )
                    cumulative += result.reward
                    reward_count += 1
                    obs = result.observation
                    if result.done:
                        break

            avg_iteration_reward = cumulative / max(1, reward_count)
            reward_history.append(round(avg_iteration_reward, 4))
            print(
                f"[train] iteration={iteration} episodes={episodes} "
                f"avg_step_reward={avg_iteration_reward:.4f}"
            )
    finally:
        await env.close()
    return rows, reward_history


def require_training_stack(*, allow_fallback: bool) -> tuple[object, object, object, object | None]:
    """Import training libraries, requiring Unsloth unless fallback is allowed."""
    try:
        from unsloth import FastLanguageModel
    except ImportError as exc:
        if not allow_fallback:
            raise RuntimeError(
                "Unsloth is required for this run but was not found.\n"
                "Install dependencies in Colab Step 2 and restart runtime, then rerun.\n"
                "If you intentionally want the transformers+peft fallback, run with --allow-fallback."
            ) from exc
        FastLanguageModel = None
        print("[train] unsloth not found; using transformers+peft fallback.")

    try:
        from datasets import Dataset
        from trl import GRPOConfig, GRPOTrainer
    except ImportError as exc:
        raise RuntimeError(
            "Missing training dependencies. Install in Colab with:\n"
            'pip install "trl>=0.15" datasets peft bitsandbytes '
            "llm-blender accelerate transformers"
        ) from exc

    return Dataset, GRPOConfig, GRPOTrainer, FastLanguageModel


def write_json(path: Path, payload: object) -> None:
    """Write JSON artifact with stable UTF-8 formatting."""
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def _run_checkpoint_eval_subprocess(
    *,
    episodes_per_difficulty: int,
    checkpoint_dir: Path,
    checkpoint_base_model: str,
    output_dir: Path,
    sandbox: bool = False,
    sandbox_drill_mode: bool = False,
    sandbox_drill_seed: int | None = None,
) -> object:
    """Run trained-checkpoint evaluation in a clean process.

    Unsloth patches model internals when imported during training. Running
    checkpoint eval in a fresh Python process avoids cross-library monkey-patch
    conflicts (for example, Qwen attention apply_qkv attribute errors).
    """
    if not checkpoint_dir.exists():
        raise FileNotFoundError(f"Adapter directory not found: {checkpoint_dir}")
    print(f"[train] launching checkpoint evaluation subprocess for {checkpoint_dir}")

    eval_output_dir = output_dir / "checkpoint_eval"
    eval_cmd = [
        sys.executable,
        str(Path(__file__).with_name("evaluate.py")),
        "--policy",
        "trained_checkpoint",
        "--episodes-per-difficulty",
        str(episodes_per_difficulty),
        "--checkpoint-dir",
        str(checkpoint_dir),
        "--checkpoint-base-model",
        checkpoint_base_model,
        "--output-dir",
        str(eval_output_dir),
    ]
    if sandbox:
        eval_cmd.append("--sandbox")
    if sandbox_drill_mode:
        eval_cmd.append("--sandbox-drill-mode")
    if sandbox_drill_seed is not None:
        eval_cmd.extend(["--sandbox-drill-seed", str(sandbox_drill_seed)])
    result = subprocess.run(
        eval_cmd,
        capture_output=True,
        text=True,
    )
    if result.stdout:
        for line in result.stdout.strip().splitlines()[-10:]:
            print(f"  [checkpoint-eval stdout] {line}")
    if result.returncode != 0:
        stderr_tail = (result.stderr or "")[-500:]
        print(f"  [checkpoint-eval FAILED] exit={result.returncode}")
        if stderr_tail:
            print(f"  [checkpoint-eval stderr] ...{stderr_tail}")
        raise RuntimeError(
            f"Checkpoint eval subprocess failed (exit={result.returncode})"
        )

    report_path = eval_output_dir / "trained_report.json"
    if not report_path.exists():
        raise FileNotFoundError(f"Missing checkpoint evaluation report: {report_path}")

    from evaluate import EvaluationReport

    payload = json.loads(report_path.read_text(encoding="utf-8"))
    report = EvaluationReport(**payload)
    print(f"[train] checkpoint eval complete: policy_used={report.policy_used} "
          f"avg_norm={report.avg_normalized_reward:.3f}")
    return report


def _seed_everything(seed: int) -> None:
    """Seed common RNGs for reproducible trajectory collection and training."""
    random.seed(seed)
    try:
        import numpy as np

        np.random.seed(seed)
    except ImportError:
        pass
    try:
        import torch

        torch.manual_seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(seed)
    except ImportError:
        pass


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Train EICC policy with GRPO.")
    parser.add_argument("--iterations", type=int, default=20)
    parser.add_argument("--episodes", type=int, default=30)
    parser.add_argument("--k", type=int, default=4, help="GRPO num_generations")
    parser.add_argument("--eval-episodes", type=int, default=5)
    parser.add_argument("--output-dir", default="artifacts/train")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--max-completion-length",
        type=int,
        default=256,
        help="Max new tokens the policy generates per action (kept short for JSON).",
    )
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument(
        "--allow-fallback",
        action="store_true",
        help="Allow transformers+peft fallback when Unsloth is unavailable.",
    )
    parser.add_argument(
        "--sandbox-drill-eval",
        action="store_true",
        help=(
            "Optional post-training sandbox drill evaluation. "
            "Requires local sandbox cluster + OPENENV_SANDBOX_* connectivity."
        ),
    )
    parser.add_argument(
        "--sandbox-drill-seed",
        type=int,
        default=None,
        help="Deterministic seed override for sandbox drill schedule.",
    )
    return parser


def _build_grpo_config(
    *,
    GRPOConfig: object,
    output_dir: Path,
    k: int,
    max_completion_length: int,
) -> object:
    """Construct GRPOConfig with TRL-version-compatible argument names."""
    params = inspect.signature(GRPOConfig).parameters
    kwargs: dict[str, object] = {
        "output_dir": str(output_dir / "grpo_output"),
        "num_train_epochs": 1,
        # More conservative LR improves PPO/GRPO stability when clip stats are saturated.
        "learning_rate": 2e-6,
        "logging_steps": 5,
        "save_steps": 100,
        "warmup_steps": 10,
    }

    # TRL naming drift across releases.
    if "num_generations" in params:
        kwargs["num_generations"] = k
    elif "num_generation" in params:
        kwargs["num_generation"] = k

    if "max_new_tokens" in params:
        kwargs["max_new_tokens"] = max_completion_length
    elif "max_completion_length" in params:
        kwargs["max_completion_length"] = max_completion_length
    elif "response_length" in params:
        kwargs["response_length"] = max_completion_length

    # GRPO requires per_device_train_batch_size to be a multiple of
    # num_generations. Setting them equal gives one prompt per device;
    # gradient_accumulation_steps controls effective batch size.
    if "per_device_train_batch_size" in params:
        kwargs["per_device_train_batch_size"] = k
    elif "train_batch_size" in params:
        kwargs["train_batch_size"] = k

    if "gradient_accumulation_steps" in params:
        kwargs["gradient_accumulation_steps"] = 2

    # Encourage diverse generations so GRPO gets meaningful reward variance.
    if "temperature" in params:
        kwargs["temperature"] = 0.8
    if "top_p" in params:
        kwargs["top_p"] = 0.95

    # Stop generation after a newline so the model can produce short,
    # clean single-JSON outputs instead of padding to max_completion_length.
    if "stop_strings" in params:
        kwargs["stop_strings"] = ["\n"]

    return GRPOConfig(**kwargs)


def main() -> None:
    """Run trajectory collection, optional GRPO training, and evaluation."""
    args = _build_parser().parse_args()
    output_dir = Path(args.output_dir)
    _seed_everything(args.seed)

    trajectories, reward_history = asyncio.run(
        collect_trajectories(iterations=args.iterations, episodes=args.episodes)
    )
    write_json(
        output_dir / "trajectories.json",
        [asdict(row) for row in trajectories],
    )
    write_json(output_dir / "reward_history.json", reward_history)

    if args.dry_run:
        print(
            f"[dry-run] collected_rows={len(trajectories)} "
            f"iterations={args.iterations} episodes={args.episodes}"
        )
        return

    Dataset, GRPOConfig, GRPOTrainer, FastLanguageModel = require_training_stack(
        allow_fallback=args.allow_fallback
    )
    dataset_rows = [
        {"prompt": row.prompt, "completion": row.completion, "reward": row.reward}
        for row in trajectories
    ]
    dataset = Dataset.from_list(dataset_rows)
    # Action-keyed lookup: the environment reward recorded when this
    # action_type was executed from this prompt during trajectory collection.
    # Much more forgiving than full (prompt, completion) string match.
    action_reward_lookup: dict[tuple[str, str], list[float]] = {}
    for row in trajectories:
        try:
            recorded_action = json.loads(row.completion)
        except json.JSONDecodeError:
            continue
        if not isinstance(recorded_action, dict):
            continue
        action_type = str(recorded_action.get("action_type", "")).strip()
        if not action_type:
            continue
        action_reward_lookup.setdefault((row.prompt, action_type), []).append(row.reward)
    trajectory_action_reward: dict[tuple[str, str], float] = {
        key: sum(values) / len(values) for key, values in action_reward_lookup.items()
    }

    model_name = "Qwen/Qwen2.5-3B-Instruct"
    target_modules = [
        "q_proj",
        "k_proj",
        "v_proj",
        "o_proj",
        "gate_proj",
        "up_proj",
        "down_proj",
    ]

    if FastLanguageModel is not None:
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name,
            max_seq_length=4096,
            load_in_4bit=True,
            dtype=None,
        )
        model = FastLanguageModel.get_peft_model(
            model,
            r=16,
            lora_alpha=16,
            lora_dropout=0,
            target_modules=target_modules,
        )
    else:
        from peft import LoraConfig, get_peft_model
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
        bnb_cfg = BitsAndBytesConfig(load_in_4bit=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map="auto",
            quantization_config=bnb_cfg,
        )
        peft_cfg = LoraConfig(
            r=16,
            lora_alpha=16,
            lora_dropout=0.0,
            target_modules=target_modules,
            bias="none",
            task_type="CAUSAL_LM",
        )
        model = get_peft_model(model, peft_cfg)

    preferred_actions: dict[str, set[str]] = {
        "triage": {"check_monitoring", "query_kb", "classify"},
        "investigation": {"check_monitoring", "probe_service", "fetch_logs", "query_kb"},
        "response": {"check_policy", "apply_fix", "notify_stakeholders", "respond"},
        "resolution": {"verify_fix", "write_postmortem", "update_kb", "resolve"},
    }
    required_fields: dict[str, tuple[str, ...]] = {
        "classify": ("category", "priority"),
        "probe_service": ("service_name", "check_type"),
        "fetch_logs": ("service_name", "time_range"),
        "respond": ("response_text", "tone"),
        "apply_fix": ("service_name", "fix_type"),
        "resolve": ("resolution_summary",),
        "notify_stakeholders": ("stakeholder", "urgency"),
        "write_postmortem": ("summary", "root_cause_description"),
        "update_kb": ("article_title", "content"),
    }

    reward_stats = {
        "total": 0,
        "valid_json": 0,
        "valid_action": 0,
        "available": 0,
        "unterminated": 0,
        "multi_json": 0,
        "extra_text": 0,
        "cap_hit": 0,
    }
    length_stats = {"sum": 0.0, "count": 0, "max": 0}

    def _prompt_field(prompt: str, key: str) -> str:
        prefix = f"{key}="
        for line in prompt.splitlines():
            if line.startswith(prefix):
                return line[len(prefix) :].strip()
        return ""

    def _mentions_known_action(text: str) -> bool:
        for action_name in KNOWN_ACTION_TYPES:
            if f'"{action_name}"' in text or f"'{action_name}'" in text:
                return True
        return False

    def _looks_terminated_json(text: str) -> bool:
        stripped = text.strip()
        return stripped.endswith("}")

    def _score_single(prompt: str, completion: str) -> float:
        reward_stats["total"] += 1
        length_stats["sum"] += len(completion)
        length_stats["count"] += 1
        if len(completion) > length_stats["max"]:
            length_stats["max"] = len(completion)

        action_payload = _extract_first_json_object(completion)
        if action_payload is None:
            if not _looks_terminated_json(completion):
                reward_stats["unterminated"] += 1
                return -0.15
            if _mentions_known_action(completion):
                return -0.08
            return -0.06

        reward_stats["valid_json"] += 1
        action_type = str(action_payload.get("action_type", "")).strip()
        if not action_type:
            return -0.02

        reward_stats["valid_action"] += 1

        available_actions_raw = _prompt_field(prompt, "available_actions")
        available_actions: set[str] = set()
        if available_actions_raw:
            try:
                parsed_available = ast.literal_eval(available_actions_raw)
                if isinstance(parsed_available, list):
                    available_actions = {str(item) for item in parsed_available}
            except (ValueError, SyntaxError):
                available_actions = set()

        phase = _prompt_field(prompt, "phase")

        score = 0.12
        if action_type in KNOWN_ACTION_TYPES:
            score += 0.08

        if action_type in available_actions:
            score += 0.24
            reward_stats["available"] += 1
        else:
            score -= 0.02

        if action_type in preferred_actions.get(phase, set()):
            score += 0.10

        needed = required_fields.get(action_type, ())
        if needed:
            present = sum(
                1 for field_name in needed if action_payload.get(field_name) not in (None, "")
            )
            score += 0.12 * (present / max(1, len(needed)))

        trajectory_reward = trajectory_action_reward.get((prompt, action_type), 0.0)
        score += 0.25 * float(trajectory_reward)

        # Strongly prefer exactly one JSON object and no trailing/leading prose.
        # These penalties must dominate so GRPO learns to produce clean outputs.
        matches = _extract_json_object_matches(completion)
        is_clean = False
        if matches:
            first = matches[0]
            prefix = completion[: first.start()].strip()
            suffix = completion[first.end() :].strip()
            if len(matches) > 1:
                reward_stats["multi_json"] += 1
                score -= 0.40  # Harsh: must learn single-object output
            if prefix or suffix:
                reward_stats["extra_text"] += 1
                score -= 0.30  # Harsh: must learn no extra text
            if len(matches) == 1 and not prefix and not suffix:
                is_clean = True
                score += 0.15  # Strong bonus for exactly one clean JSON object

        completion_len = len(completion)
        cap_threshold = max(16, int(args.max_completion_length * 0.95))
        if completion_len >= cap_threshold:
            reward_stats["cap_hit"] += 1
            score -= 0.25  # Severe: model must learn to stop early
        if is_clean and completion_len <= 120:
            score += 0.05  # Reward concise clean completions
        elif completion_len > 400:
            score -= 0.20
        elif completion_len > 240:
            score -= 0.12
        elif completion_len > 160:
            score -= 0.06

        return max(-1.0, min(1.0, score))

    def reward_function(
        prompts: list[str],
        completions: list[str],
        **_: object,
    ) -> list[float]:
        rewards: list[float] = []
        for raw_prompt, raw_completion in zip(prompts, completions):
            prompt = raw_prompt if isinstance(raw_prompt, str) else str(raw_prompt or "")
            completion = (
                raw_completion if isinstance(raw_completion, str) else str(raw_completion or "")
            )
            try:
                rewards.append(_score_single(prompt, completion))
            except Exception as exc:  # pragma: no cover - defensive guard
                # Never let a single bad sample crash the reward callback.
                print(f"[reward] scoring error: {exc!r}; sample length={len(completion)}")
                rewards.append(-0.10)

        # Keep logs concise: report health checkpoints and only print
        # sample completions when signal quality is poor.
        batch_count = reward_function.__dict__.setdefault("_call_count", 0) + 1
        reward_function.__dict__["_call_count"] = batch_count
        if batch_count == 1 or batch_count % 25 == 0:
            total = max(1, reward_stats["total"])
            samples = max(1, length_stats["count"])
            avg_len = length_stats["sum"] / samples
            valid_json_rate = reward_stats["valid_json"] / total
            valid_action_rate = reward_stats["valid_action"] / total
            available_rate = reward_stats["available"] / total
            unterminated_rate = reward_stats["unterminated"] / total
            cap_hit_rate = reward_stats["cap_hit"] / total
            multi_json_rate = reward_stats["multi_json"] / total
            extra_text_rate = reward_stats["extra_text"] / total
            print(
                "[reward] batch={} total={} valid_json={:.0%} valid_action={:.0%} "
                "available={:.0%} unterminated={:.0%} cap_hit={:.0%} "
                "multi_json={:.0%} extra_text={:.0%} avg_len={:.0f} max_len={}".format(
                    batch_count,
                    reward_stats["total"],
                    valid_json_rate,
                    valid_action_rate,
                    available_rate,
                    unterminated_rate,
                    cap_hit_rate,
                    multi_json_rate,
                    extra_text_rate,
                    avg_len,
                    length_stats["max"],
                )
            )
            unhealthy_signal = (
                valid_json_rate < 0.90
                or valid_action_rate < 0.80
                or unterminated_rate > 0.10
                or cap_hit_rate > 0.30
                or multi_json_rate > 0.10
                or extra_text_rate > 0.10
            )
            if unhealthy_signal:
                print(
                    "[reward] warning: noisy outputs detected; prefer exactly one compact JSON object."
                )
            if unhealthy_signal and completions:
                preview_raw = completions[0]
                preview = preview_raw if isinstance(preview_raw, str) else str(preview_raw or "")
                preview = preview.replace("\n", " ")
                if len(preview) > 160:
                    preview = preview[:160] + "..."
                print(f"[reward] sample completion: {preview}")
        return rewards

    config = _build_grpo_config(
        GRPOConfig=GRPOConfig,
        output_dir=output_dir,
        k=args.k,
        max_completion_length=args.max_completion_length,
    )
    trainer_params = inspect.signature(GRPOTrainer).parameters
    trainer_kwargs: dict[str, object] = {"model": model}

    # TRL / Unsloth naming drift: some versions use `config`, others `args`.
    if "config" in trainer_params:
        trainer_kwargs["config"] = config
    elif "args" in trainer_params:
        trainer_kwargs["args"] = config

    # Reward callback naming is stable in recent TRL but keep compatibility.
    if "reward_funcs" in trainer_params:
        trainer_kwargs["reward_funcs"] = [reward_function]
    elif "reward_function" in trainer_params:
        trainer_kwargs["reward_function"] = reward_function

    if "train_dataset" in trainer_params:
        trainer_kwargs["train_dataset"] = dataset
    elif "dataset" in trainer_params:
        trainer_kwargs["dataset"] = dataset

    if "tokenizer" in trainer_params:
        trainer_kwargs["tokenizer"] = tokenizer
    elif "processing_class" in trainer_params:
        trainer_kwargs["processing_class"] = tokenizer
    trainer = GRPOTrainer(**trainer_kwargs)
    trainer.train()
    trainer.save_model(str(output_dir / "trained_adapter"))

    baseline = evaluate_policy(
        policy="baseline",
        episodes_per_difficulty=args.eval_episodes,
    )
    trained_adapter_dir = output_dir / "trained_adapter"
    try:
        trained = _run_checkpoint_eval_subprocess(
            episodes_per_difficulty=args.eval_episodes,
            checkpoint_dir=trained_adapter_dir,
            checkpoint_base_model=model_name,
            output_dir=output_dir,
        )
    except Exception as exc:
        print(
            "[train] checkpoint-based evaluation unavailable; "
            f"falling back to trained_heuristic policy ({exc})."
        )
        trained = evaluate_policy(
            policy="trained_heuristic",
            episodes_per_difficulty=args.eval_episodes,
        )
    # Structured, truthful behavior diff derived from actual eval metrics.
    trained.behavior_examples = behavior_diffs(baseline, trained)
    write_json(output_dir / "baseline_report.json", asdict(baseline))
    write_json(output_dir / "trained_report.json", asdict(trained))
    plot_reports(baseline, trained, output_dir)
    trained.print_comparison(baseline)
    print(f"[train] trained_report.policy_used={trained.policy_used}")

    if args.sandbox_drill_eval:
        print("[train] running optional sandbox drill evaluation (add-on)...")
        sbx_baseline = evaluate_policy(
            policy="baseline",
            episodes_per_difficulty=args.eval_episodes,
            sandbox=True,
            sandbox_drill_mode=True,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        try:
            sbx_trained = _run_checkpoint_eval_subprocess(
                episodes_per_difficulty=args.eval_episodes,
                checkpoint_dir=trained_adapter_dir,
                checkpoint_base_model=model_name,
                output_dir=output_dir / "sandbox_drill_eval",
                sandbox=True,
                sandbox_drill_mode=True,
                sandbox_drill_seed=args.sandbox_drill_seed,
            )
        except Exception as exc:
            print(
                "[train] sandbox checkpoint eval unavailable; "
                f"falling back to trained_heuristic policy ({exc})."
            )
            sbx_trained = evaluate_policy(
                policy="trained_heuristic",
                episodes_per_difficulty=args.eval_episodes,
                sandbox=True,
                sandbox_drill_mode=True,
                sandbox_drill_seed=args.sandbox_drill_seed,
            )
        sbx_trained.behavior_examples = behavior_diffs(sbx_baseline, sbx_trained)
        write_json(output_dir / "baseline_sandbox_drill_report.json", asdict(sbx_baseline))
        write_json(output_dir / "trained_sandbox_drill_report.json", asdict(sbx_trained))
        print("[train] sandbox drill reports written.")


if __name__ == "__main__":
    main()