eval_baseline.py

"""Pre-training evaluation harness.

Runs a base LLM agent against N episodes of DriftEnv and reports:
  - mean episode reward
  - mean fraction of max (max = 30.0 per episode)
  - drift-sensitive-step accuracy  <-- the KEY metric
  - breakdown by drift event
  - sample trajectories

The whole point is to confirm the base model's drift-sensitive-step accuracy
sits in the 20%-40% training-headroom zone. Above 60% means our task is too
easy and training won't show improvement. Below 10% means our task is too
hard for any RL within the time budget.

Usage:
    API_BASE_URL=https://api.groq.com/openai/v1 \\
    HF_TOKEN=gsk_... \\
    MODEL_NAME=llama-3.1-8b-instant \\
    PYTHONPATH=. python3 eval_baseline.py --episodes 10
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import defaultdict
from typing import Any

from dotenv import load_dotenv

from drift_env.environment import DriftEnv
from drift_env.llm_agent import LLMAgent
from drift_env.episodes import generate_episode
from drift_env.policy import drift_direction

load_dotenv()


def run_one_episode(
    agent: LLMAgent, seed: int, verbose: bool = False,
) -> dict[str, Any]:
    env = DriftEnv()
    obs = env.reset(seed=seed, episode_id=f"eval_{seed}")
    ep_plan = generate_episode(seed=seed, episode_id=f"eval_{seed}")

    total_reward = 0.0
    breakdown_totals = defaultdict(float)
    drift_sensitive_total = 0
    drift_sensitive_correct = 0
    per_drift: dict[str, dict[str, int]] = defaultdict(
        lambda: {"total": 0, "correct": 0}
    )
    per_direction: dict[str, dict[str, int]] = defaultdict(
        lambda: {"total": 0, "correct": 0}
    )
    trajectory = []

    for i, step_plan in enumerate(ep_plan.steps):
        action, raw = agent.act(obs)
        result = env.step(action)

        total_reward += result.reward
        for k, v in result.info["breakdown"].items():
            breakdown_totals[k] += v

        sensitive_to = step_plan.drift_sensitive_to
        is_correct = result.info["breakdown"]["compliance"] >= 1.0
        if sensitive_to is not None:
            drift_sensitive_total += 1
            per_drift[sensitive_to]["total"] += 1
            direction = drift_direction(sensitive_to)
            if direction is not None:
                per_direction[direction]["total"] += 1
            if is_correct:
                drift_sensitive_correct += 1
                per_drift[sensitive_to]["correct"] += 1
                if direction is not None:
                    per_direction[direction]["correct"] += 1

        trajectory.append({
            "step": i,
            "email_kind": step_plan.email.kind.value,
            "email_id": step_plan.email.id,
            "action": action.model_dump(exclude_none=True),
            "correct_hint": step_plan.correct_action_hint,
            "drift_sensitive_to": sensitive_to,
            "reward": result.reward,
            "compliance": result.info["breakdown"]["compliance"],
        })

        if verbose:
            sens = f" [SENS→{sensitive_to}]" if sensitive_to else ""
            print(f"  step {i:>2} [{step_plan.email.kind.value:<8}] "
                  f"reward={result.reward:.2f} comp={result.info['breakdown']['compliance']:.2f}{sens}")

        if result.done:
            break
        if result.observation is not None:
            obs = result.observation

    return {
        "seed": seed,
        "total_reward": round(total_reward, 4),
        "max_reward": 30.0,
        "frac_of_max": round(total_reward / 30.0, 4),
        "breakdown_totals": {k: round(v, 4) for k, v in breakdown_totals.items()},
        "drift_sensitive_total": drift_sensitive_total,
        "drift_sensitive_correct": drift_sensitive_correct,
        "drift_sensitive_acc": (
            round(drift_sensitive_correct / drift_sensitive_total, 4)
            if drift_sensitive_total else None
        ),
        "per_drift": {k: dict(v) for k, v in per_drift.items()},
        "per_direction": {k: dict(v) for k, v in per_direction.items()},
        "trajectory": trajectory,
    }


def summarise(results: list[dict[str, Any]]) -> dict[str, Any]:
    n = len(results)
    mean_reward = sum(r["total_reward"] for r in results) / n
    mean_frac = sum(r["frac_of_max"] for r in results) / n
    dst = sum(r["drift_sensitive_total"] for r in results)
    dsc = sum(r["drift_sensitive_correct"] for r in results)

    per_drift_agg: dict[str, dict[str, int]] = defaultdict(
        lambda: {"total": 0, "correct": 0}
    )
    per_direction_agg: dict[str, dict[str, int]] = defaultdict(
        lambda: {"total": 0, "correct": 0}
    )
    for r in results:
        for name, stats in r["per_drift"].items():
            per_drift_agg[name]["total"] += stats["total"]
            per_drift_agg[name]["correct"] += stats["correct"]
        for direction, stats in r.get("per_direction", {}).items():
            per_direction_agg[direction]["total"] += stats["total"]
            per_direction_agg[direction]["correct"] += stats["correct"]

    return {
        "episodes": n,
        "mean_reward": round(mean_reward, 4),
        "mean_frac_of_max": round(mean_frac, 4),
        "drift_sensitive_total": dst,
        "drift_sensitive_correct": dsc,
        "drift_sensitive_acc": round(dsc / dst, 4) if dst else None,
        "per_direction": {
            k: {
                **v,
                "acc": round(v["correct"] / v["total"], 4) if v["total"] else None,
            }
            for k, v in per_direction_agg.items()
        },
        "per_drift": {
            k: {
                **v,
                "acc": round(v["correct"] / v["total"], 4) if v["total"] else None,
            }
            for k, v in per_drift_agg.items()
        },
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--episodes", type=int, default=10)
    ap.add_argument("--start-seed", type=int, default=100)
    ap.add_argument("--verbose", action="store_true")
    ap.add_argument("--save", type=str, default="eval_results.json")
    args = ap.parse_args()

    api_base = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
    api_key = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
    model = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")

    if not api_key:
        print("ERROR: set HF_TOKEN (or API_KEY) env var.", file=sys.stderr)
        return 1

    print(f"Model: {model}")
    print(f"Base URL: {api_base}")
    print(f"Episodes: {args.episodes}")
    print("-" * 60)

    agent = LLMAgent(api_key=api_key, base_url=api_base, model=model)

    results = []
    t0 = time.time()
    for i in range(args.episodes):
        seed = args.start_seed + i
        if args.verbose:
            print(f"\n=== Episode seed={seed} ===")
        try:
            r = run_one_episode(agent, seed=seed, verbose=args.verbose)
        except Exception as e:
            print(f"  FAILED: {e}")
            continue
        results.append(r)
        dsa = r["drift_sensitive_acc"]
        dsa_str = f"{dsa:.2%}" if dsa is not None else "n/a"
        pd = r.get("per_direction", {})
        def _fmt(dir_):
            s = pd.get(dir_, {})
            t = s.get("total", 0); c = s.get("correct", 0)
            return f"{dir_[:4]}={c}/{t}"
        per_dir_str = "  ".join(_fmt(d) for d in ("tightening", "loosening", "neutral"))
        print(f"  seed={seed} reward={r['total_reward']:.2f}/30 "
              f"drift_acc={dsa_str} "
              f"({r['drift_sensitive_correct']}/{r['drift_sensitive_total']})  "
              f"[{per_dir_str}]")

    dt = time.time() - t0
    summary = summarise(results)
    print("\n" + "=" * 60)
    print("SUMMARY")
    print("=" * 60)
    print(json.dumps(summary, indent=2))
    print(f"\nTook {dt:.1f}s for {args.episodes} episodes.")

    # Interpretation
    dsa = summary["drift_sensitive_acc"]
    if dsa is None:
        print("\n[warn] no drift-sensitive steps encountered.")
    elif dsa < 0.10:
        print(f"\n[warn] drift-sensitive acc = {dsa:.0%}. Too hard for RL — redesign.")
    elif dsa > 0.60:
        print(f"\n[warn] drift-sensitive acc = {dsa:.0%}. Too easy — base model already solves it.")
    elif dsa > 0.40:
        print(f"\n[ok-ish] drift-sensitive acc = {dsa:.0%}. In the training zone but on the easy side.")
    else:
        print(f"\n[OK] drift-sensitive acc = {dsa:.0%}. In the sweet spot for training.")

    with open(args.save, "w") as f:
        json.dump({"summary": summary, "results": results}, f, indent=2)
    print(f"\nFull results written to {args.save}")
    return 0


if __name__ == "__main__":
    sys.exit(main())