evaluate.py

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

This script runs deterministic incident-mode episodes and reports formal metrics:
- normalized reward
- SLA compliance
- tool efficiency
- root-cause accuracy proxy
- long-horizon consistency
- 8 tracked behavioral skills

Usage examples:
    python evaluate.py --policy baseline --episodes-per-difficulty 3
    python evaluate.py --policy compare --episodes-per-difficulty 5 --plot
"""

from __future__ import annotations

import argparse
import asyncio
import json
import re
from collections import defaultdict
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Callable, Literal

from env.environment import CustomerSupportEnv
from models.action import ActionAdapter
from models.observation import Observation
from sandbox.env_adapter import SandboxEnv

PolicyKind = Literal["baseline", "trained", "trained_heuristic", "trained_checkpoint"]

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

DIFFICULTIES: tuple[str, ...] = ("easy", "medium", "hard", "nightmare")
TOOL_ACTIONS = frozenset(
    [
        "check_monitoring",
        "probe_service",
        "fetch_logs",
        "fetch_user_data",
        "check_billing",
        "query_kb",
        "check_policy",
        "query_incident_history",
        "follow_runbook_step",
    ]
)
POLICY_SENSITIVE_ACTIONS = frozenset(
    ["apply_fix", "escalate", "notify_stakeholders", "update_kb", "resolve"]
)
TRACKED_SKILLS: tuple[str, ...] = (
    "investigation_before_action",
    "kb_cross_verification",
    "policy_checking",
    "stakeholder_proactivity",
    "root_cause_accuracy",
    "tone_matching",
    "resource_efficiency",
    "red_herring_dismissal",
)

REQUIRED_FIELDS_BY_ACTION: dict[str, tuple[str, ...]] = {
    "classify": ("category", "priority"),
    "route": ("department",),
    "probe_service": ("service_name", "check_type"),
    "fetch_logs": ("service_name", "time_range"),
    "fetch_user_data": ("customer_id",),
    "check_billing": ("customer_id",),
    "query_kb": ("query",),
    "check_policy": ("policy_type",),
    "query_incident_history": ("query",),
    "respond": ("response_text", "tone"),
    "apply_fix": ("service_name", "fix_type"),
    "verify_fix": ("service_name",),
    "rollback_fix": ("service_name",),
    "request_info": ("question_to_customer",),
    "escalate": ("reason", "target_team"),
    "resolve": ("resolution_summary",),
    "notify_stakeholders": ("stakeholder", "urgency", "message"),
    "follow_runbook_step": ("runbook_id", "step_index"),
    "write_postmortem": ("summary", "root_cause_description"),
    "update_kb": ("article_title", "content"),
}

# Literal value whitelists per action type (sourced from models/action.py).
# These keep checkpoint hallucinations (e.g. category="cyber") from reaching
# the env. If a value is outside the whitelist, sanitization will drop the
# candidate and use the deterministic heuristic fallback for that step.
_LITERAL_FIELDS_BY_ACTION: dict[str, dict[str, frozenset[str]]] = {
    "classify": {
        "category": frozenset(
            {
                "billing",
                "bug_report",
                "feature_request",
                "account_access",
                "general_inquiry",
                "cancellation",
            }
        ),
        "priority": frozenset({"low", "medium", "high", "critical"}),
    },
    "route": {
        "department": frozenset({"billing", "technical", "account", "general"}),
    },
    "respond": {
        "tone": frozenset({"formal", "empathetic", "concise"}),
    },
    "escalate": {
        "target_team": frozenset({"l2_support", "engineering", "management"}),
    },
    "probe_service": {
        "check_type": frozenset({"logs", "resources", "connections", "config"}),
    },
    "fetch_logs": {
        "time_range": frozenset({"last_5m", "last_15m", "last_1h"}),
    },
    "check_policy": {
        "policy_type": frozenset(
            {"refund", "escalation", "sla", "compensation", "communication"}
        ),
    },
    "notify_stakeholders": {
        "stakeholder": frozenset(
            {"vp_engineering", "legal", "support_lead", "all"}
        ),
        "urgency": frozenset({"info", "warning", "critical"}),
    },
}

PHASE_ALLOWED_ACTIONS: dict[str, set[str]] = {
    "triage": {"check_monitoring", "query_kb", "classify"},
    "investigation": {
        "check_monitoring",
        "query_kb",
        "fetch_logs",
        "probe_service",
        "route",
        "fetch_user_data",
    },
    "response": {"check_policy", "notify_stakeholders", "apply_fix", "respond"},
    "resolution": {"verify_fix", "write_postmortem", "update_kb", "resolve", "respond"},
}


@dataclass(slots=True)
class EpisodeReport:
    """Per-episode evaluation record."""

    difficulty: str
    normalized_reward: float
    raw_reward: float
    sla_compliant: float
    tool_efficiency: float
    root_cause_accuracy: float
    long_horizon_consistency: float
    skill_scores: dict[str, float]
    actions: list[str]


@dataclass
class EvaluationReport:
    """Structured output from evaluate.py — used for demo and blog."""

    avg_normalized_reward: float
    avg_raw_reward: float
    sla_compliance_rate: float
    tool_efficiency: float
    root_cause_accuracy: float
    long_horizon_consistency: float
    skill_scores: dict[str, float]
    per_difficulty: dict[str, float]
    reward_history: list[float]
    raw_reward_history: list[float]
    per_difficulty_reward_history: dict[str, list[float]] = field(default_factory=dict)
    behavior_examples: list[str] = field(default_factory=list)
    policy_used: str = "unknown"
    episodes_per_difficulty: int = 0

    def print_comparison(self, baseline: "EvaluationReport") -> None:
        """Pretty-print before vs after for demo."""
        print("=" * 60)
        print("BEFORE TRAINING -> AFTER TRAINING")
        print("=" * 60)
        mult = self.avg_normalized_reward / max(0.0001, baseline.avg_normalized_reward)
        print(
            f"Normalized Reward:  {baseline.avg_normalized_reward:.3f} -> "
            f"{self.avg_normalized_reward:.3f}  ({mult:.1f}x)"
        )
        print(
            f"Raw Cumulative:     {baseline.avg_raw_reward:+.3f} -> "
            f"{self.avg_raw_reward:+.3f}"
        )
        print(
            f"SLA Compliance:     {baseline.sla_compliance_rate:.0%} -> "
            f"{self.sla_compliance_rate:.0%}"
        )
        print(
            f"Root Cause Accuracy:{baseline.root_cause_accuracy:.0%} -> "
            f"{self.root_cause_accuracy:.0%}"
        )
        print(
            f"Tool Efficiency:    {baseline.tool_efficiency:.2f} -> "
            f"{self.tool_efficiency:.2f}"
        )
        print(
            "Long-Horizon Consistency: "
            f"{baseline.long_horizon_consistency:.2f} -> {self.long_horizon_consistency:.2f}"
        )
        for skill in TRACKED_SKILLS:
            base_score = baseline.skill_scores.get(skill, 0.0)
            score = self.skill_scores.get(skill, 0.0)
            print(f"  {skill}: {base_score:.0%} -> {score:.0%}")


@dataclass
class TransferReport:
    """Cross-backend transfer summary (simulated -> sandbox)."""

    trained_policy: str
    episodes_per_difficulty: int
    simulated: dict[str, Any]
    sandbox: dict[str, Any]
    transfer: dict[str, Any]


@dataclass(slots=True)
class PolicyState:
    """Minimal per-episode policy memory for deterministic action selection."""

    has_checked_monitoring: bool = False
    has_queried_kb: bool = False
    has_checked_policy: bool = False
    has_notified: bool = False
    has_applied_fix: bool = False
    has_verified_fix: bool = False
    has_resolved: bool = False
    has_fetched_logs: bool = False
    has_probed_service: bool = False
    has_written_postmortem: bool = False
    has_updated_kb: bool = False
    known_service: str = "auth"
    known_fix_type: str = "restart_service"
    known_root_cause: str = ""


def _is_trained_policy(policy: PolicyKind) -> bool:
    return policy in ("trained", "trained_heuristic")


def _priority_from_max_steps(max_steps: int) -> str:
    if max_steps >= 80:
        return "critical"
    if max_steps >= 70:
        return "critical"
    if max_steps >= 50:
        return "high"
    return "medium"


def _default_customer_id(obs: Observation) -> str:
    if obs.ticket_id.startswith("CUST-"):
        return obs.ticket_id
    return "CUST-001"


def _fallback_action(obs: Observation) -> dict[str, object]:
    phase = obs.incident_phase or "investigation"
    if phase in ("triage", "investigation"):
        return {"action_type": "check_monitoring", "service_name": None}
    if phase == "response":
        return {
            "action_type": "respond",
            "response_text": "We are investigating and will share updates shortly.",
            "tone": "empathetic",
        }
    return {
        "action_type": "resolve",
        "resolution_summary": "Incident reviewed and currently stable.",
        "offered_compensation": None,
    }


def _pick_impacted_service(obs: Observation, default_service: str) -> str:
    status = obs.system_status or {}
    for service_name in sorted(status.keys()):
        if status[service_name] == "down":
            return service_name
    for service_name in sorted(status.keys()):
        if status[service_name] == "degraded":
            return service_name
    return default_service


# Maps root_cause failure modes to the correct fix_type.
_FAILURE_FIX_MAP: dict[str, str] = {
    "rate_limiting": "restart_service",
    "token_expiry": "config_change",
    "config_corruption": "config_change",
    "oom": "memory_increase",
    "connection_pool_exhaustion": "restart_service",
    "replication_lag": "schema_migration",
    "gateway_timeout": "restart_service",
    "validation_errors": "config_change",
    "idempotency_failure": "data_fix",
    "batch_job_runaway": "memory_increase",
    "query_timeout": "schema_migration",
    "stale_cache": "restart_service",
    "queue_overflow": "restart_service",
    "template_error": "config_change",
    "rate_exceeded": "memory_increase",
}


def _extract_root_cause_from_facts(
    obs: Observation, state: PolicyState
) -> None:
    """Update state with root cause and fix info from probe/log evidence."""
    facts = obs.known_facts or {}
    # Probe results contain "observed failure signature: <mode>" for high-observability services
    for key, value in facts.items():
        if not isinstance(key, str):
            continue
        if key.startswith("probe:"):
            if isinstance(value, dict):
                findings = value.get("findings", [])
                if isinstance(findings, list):
                    for f in findings:
                        if isinstance(f, str) and "observed failure signature:" in f:
                            mode = f.split("observed failure signature:")[-1].strip()
                            if mode and mode != "unknown":
                                state.known_root_cause = mode
                                if mode in _FAILURE_FIX_MAP:
                                    state.known_fix_type = _FAILURE_FIX_MAP[mode]


def choose_policy_action(
    obs: Observation,
    state: PolicyState,
    policy: PolicyKind,
) -> dict[str, object]:
    """Deterministic heuristic policy used for baseline/trained-style evaluations."""
    available = set(obs.available_actions)
    phase = obs.incident_phase or "investigation"
    trained_mode = _is_trained_policy(policy)
    customer_id = _default_customer_id(obs)
    state.known_service = _pick_impacted_service(obs, state.known_service)

    # Trained policy extracts root cause from evidence gathered during investigation
    if trained_mode:
        _extract_root_cause_from_facts(obs, state)

    if phase == "triage":
        if "check_monitoring" in available and not state.has_checked_monitoring:
            state.has_checked_monitoring = True
            return {"action_type": "check_monitoring", "service_name": None}
        if "query_kb" in available and not state.has_queried_kb and trained_mode:
            state.has_queried_kb = True
            return {"action_type": "query_kb", "query": "service outage"}
        if "classify" in available:
            priority = _priority_from_max_steps(obs.max_steps)
            if policy == "baseline":
                priority = "medium"
            return {
                "action_type": "classify",
                "category": "bug_report",
                "priority": priority,
            }

    if phase == "investigation":
        if "check_monitoring" in available and not state.has_checked_monitoring:
            state.has_checked_monitoring = True
            return {"action_type": "check_monitoring", "service_name": None}
        if "query_kb" in available and not state.has_queried_kb:
            state.has_queried_kb = True
            return {"action_type": "query_kb", "query": "service outage"}
        # Trained policy probes the actual impacted service to get root cause
        if "probe_service" in available and not state.has_probed_service:
            state.has_probed_service = True
            service = "payments" if policy == "baseline" else state.known_service
            return {
                "action_type": "probe_service",
                "service_name": service,
                "check_type": "logs",
            }
        if trained_mode and "fetch_logs" in available and not state.has_fetched_logs:
            state.has_fetched_logs = True
            return {
                "action_type": "fetch_logs",
                "service_name": state.known_service,
                "time_range": "last_15m",
            }
        if "route" in available:
            return {"action_type": "route", "department": "technical"}
        if "fetch_user_data" in available:
            return {"action_type": "fetch_user_data", "customer_id": customer_id}

    if phase == "response":
        if trained_mode and "check_policy" in available and not state.has_checked_policy:
            state.has_checked_policy = True
            return {"action_type": "check_policy", "policy_type": "compensation"}
        if "apply_fix" in available and not state.has_applied_fix:
            state.has_applied_fix = True
            # Trained policy uses the correct service AND fix_type from evidence
            service = "payments" if policy == "baseline" else state.known_service
            fix_type = "restart_service" if policy == "baseline" else state.known_fix_type
            return {
                "action_type": "apply_fix",
                "service_name": service,
                "fix_type": fix_type,
            }
        if "notify_stakeholders" in available and not state.has_notified and trained_mode:
            state.has_notified = True
            return {
                "action_type": "notify_stakeholders",
                "stakeholder": "all",
                "message": "Incident impact assessed and remediation underway.",
                "urgency": "warning",
            }
        if "respond" in available:
            tone = "formal" if policy == "baseline" else "empathetic"
            return {
                "action_type": "respond",
                "response_text": "We are actively working the incident and will provide updates.",
                "tone": tone,
            }

    if phase == "resolution":
        if "verify_fix" in available and not state.has_verified_fix:
            state.has_verified_fix = True
            return {"action_type": "verify_fix", "service_name": state.known_service}
        if "write_postmortem" in available and trained_mode and not state.has_written_postmortem:
            state.has_written_postmortem = True
            return {
                "action_type": "write_postmortem",
                "summary": "Incident resolved after service remediation and verification.",
                "root_cause_description": f"Root cause: {state.known_root_cause or 'service instability'}.",
                "remediation_steps": ["checked monitoring", "applied fix", "verified recovery"],
                "prevention_measures": ["refresh runbook", "add targeted alert"],
            }
        if "update_kb" in available and trained_mode and not state.has_updated_kb:
            state.has_updated_kb = True
            return {
                "action_type": "update_kb",
                "article_title": "Incident triage playbook",
                "content": "Verify service health and probe root cause before applying fixes.",
                "tags": ["incident", "triage"],
            }
        if "notify_stakeholders" in available and not state.has_notified and trained_mode:
            state.has_notified = True
            return {
                "action_type": "notify_stakeholders",
                "stakeholder": "all",
                "message": "Incident fully resolved. Postmortem complete.",
                "urgency": "info",
            }
        if "resolve" in available and not state.has_resolved:
            state.has_resolved = True
            return {
                "action_type": "resolve",
                "resolution_summary": "Incident resolved with verified service recovery.",
                "offered_compensation": None,
            }
        if "respond" in available:
            return {
                "action_type": "respond",
                "response_text": "Systems are stable and we are monitoring closely.",
                "tone": "empathetic",
            }
    return _fallback_action(obs)


def _extract_first_json_action(text: str) -> dict[str, object] | None:
    """Parse the first JSON object from model text output."""
    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 _build_model_prompt(obs: Observation) -> str:
    """Build inference prompt for checkpoint policy action generation."""
    parts = [
        "Return exactly ONE compact JSON object and nothing else.",
        "Pick action_type only from available_actions.",
        "Include all required fields for that action.",
        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('Example: {"action_type":"check_monitoring","service_name":null}')
    return "\n".join(parts)


def _sanitize_checkpoint_action(
    *,
    obs: Observation,
    state: PolicyState,
    payload: dict[str, object] | None,
    decoded_text: str,
) -> dict[str, object]:
    """Return a safe action for env.step, falling back to heuristic when needed.

    The sanitizer is intentionally defensive. Checkpoint outputs sometimes
    contain hallucinated literals (e.g. category="cyber"), invalid runbook IDs
    (e.g. "RB-02"), or out-of-range step indices. Any unrecoverable issue
    returns the deterministic heuristic fallback so env.step never crashes.
    """
    fallback = choose_policy_action(obs, state, "trained_heuristic")

    try:
        if not isinstance(payload, dict):
            return fallback

        # Reject obviously broken outputs (HTML tags, chat-style "Human:" bleed).
        if re.search(r"<[^>]+>", decoded_text or "") or "Human:" in (decoded_text or ""):
            return fallback

        action_type = str(payload.get("action_type", "")).strip()
        available = set(obs.available_actions)
        if not action_type or action_type not in available:
            return fallback

        # Keep only JSON-safe, env-serializable values.
        candidate: dict[str, object] = {"action_type": action_type}
        for key, value in payload.items():
            if key == "action_type":
                continue
            if isinstance(value, (str, int, float, bool)) or value is None:
                candidate[str(key)] = value
            elif isinstance(value, list):
                candidate[str(key)] = value

        # ---------- per-action repairs ----------

        if action_type == "follow_runbook_step":
            # Map legacy `step` -> canonical `step_index`.
            if candidate.get("step_index") in (None, "") and candidate.get("step") not in (None, ""):
                try:
                    candidate["step_index"] = int(candidate["step"])  # type: ignore[arg-type]
                except (TypeError, ValueError):
                    pass
            # Always anchor runbook_id to the env-provided suggested runbook
            # to avoid hallucinated IDs (e.g. "RB-02") that don't exist.
            if isinstance(obs.suggested_runbook, dict):
                suggested_id = obs.suggested_runbook.get("runbook_id")
                if isinstance(suggested_id, str) and suggested_id:
                    candidate["runbook_id"] = suggested_id
            if candidate.get("runbook_id") in (None, ""):
                return fallback
            # Keep step_index within the suggested runbook's known indices.
            if isinstance(obs.suggested_runbook, dict):
                suggested_steps = obs.suggested_runbook.get("steps")
                if isinstance(suggested_steps, list):
                    valid_indices: list[int] = []
                    for item in suggested_steps:
                        if isinstance(item, dict):
                            idx = item.get("step_index")
                            if isinstance(idx, int):
                                valid_indices.append(idx)
                    if valid_indices:
                        try:
                            current_idx = int(candidate.get("step_index"))  # type: ignore[arg-type]
                        except (TypeError, ValueError):
                            current_idx = valid_indices[0]
                        if current_idx not in valid_indices:
                            current_idx = valid_indices[0]
                        candidate["step_index"] = current_idx
                    else:
                        candidate["step_index"] = 0
                else:
                    candidate.setdefault("step_index", 0)
            else:
                candidate.setdefault("step_index", 0)
            candidate.pop("step", None)

        if action_type in {"apply_fix", "verify_fix", "rollback_fix", "probe_service", "fetch_logs"}:
            # Anchor service_name to a real service in current system_status.
            status = obs.system_status if isinstance(obs.system_status, dict) else {}
            allowed_services = set(status.keys())
            service_value = candidate.get("service_name")
            if not isinstance(service_value, str) or service_value not in allowed_services:
                impacted = _pick_impacted_service(obs, state.known_service)
                if isinstance(impacted, str) and impacted:
                    candidate["service_name"] = impacted
                else:
                    return fallback

        if action_type == "apply_fix":
            fix_type = candidate.get("fix_type")
            if not isinstance(fix_type, str) or not fix_type.strip():
                # Use deterministic mapping from known root cause to fix.
                state_fix = state.known_fix_type or "restart_service"
                candidate["fix_type"] = state_fix

        if action_type in {"fetch_user_data", "check_billing"}:
            customer_value = candidate.get("customer_id")
            if not isinstance(customer_value, str) or not customer_value.strip():
                candidate["customer_id"] = _default_customer_id(obs)

        if action_type == "query_kb":
            query_value = candidate.get("query")
            if not isinstance(query_value, str) or not query_value.strip():
                candidate["query"] = obs.incident_title or "incident"

        if action_type == "query_incident_history":
            query_value = candidate.get("query")
            if not isinstance(query_value, str) or not query_value.strip():
                candidate["query"] = obs.incident_title or "similar incidents"

        if action_type == "respond":
            response_value = candidate.get("response_text")
            if not isinstance(response_value, str) or not response_value.strip():
                candidate["response_text"] = (
                    "We are investigating and will share updates shortly."
                )

        if action_type == "request_info":
            question_value = candidate.get("question_to_customer")
            if not isinstance(question_value, str) or not question_value.strip():
                candidate["question_to_customer"] = (
                    "Could you share the affected user IDs and the exact time?"
                )

        if action_type == "escalate":
            reason_value = candidate.get("reason")
            if not isinstance(reason_value, str) or not reason_value.strip():
                candidate["reason"] = "Customer impact requires specialist support."

        if action_type == "resolve":
            summary_value = candidate.get("resolution_summary")
            if not isinstance(summary_value, str) or not summary_value.strip():
                candidate["resolution_summary"] = "Incident reviewed and currently stable."

        if action_type == "notify_stakeholders":
            message_value = candidate.get("message")
            if not isinstance(message_value, str) or not message_value.strip():
                candidate["message"] = "Status update on the active incident."

        if action_type == "write_postmortem":
            summary_value = candidate.get("summary")
            if not isinstance(summary_value, str) or not summary_value.strip():
                candidate["summary"] = "Incident resolved with documented remediation."
            cause_value = candidate.get("root_cause_description")
            if not isinstance(cause_value, str) or not cause_value.strip():
                candidate["root_cause_description"] = (
                    state.known_root_cause or "Service degradation handled."
                )

        if action_type == "update_kb":
            title_value = candidate.get("article_title")
            if not isinstance(title_value, str) or not title_value.strip():
                candidate["article_title"] = obs.incident_title or "Incident KB update"
            content_value = candidate.get("content")
            if not isinstance(content_value, str) or not content_value.strip():
                candidate["content"] = "Captured remediation steps and verification."

        # ---------- literal whitelist enforcement ----------

        literal_fields = _LITERAL_FIELDS_BY_ACTION.get(action_type, {})
        for field_name, allowed in literal_fields.items():
            value = candidate.get(field_name)
            if isinstance(value, str) and value in allowed:
                continue
            # If invalid, drop the value so a default (if any) can apply,
            # otherwise fall back. We avoid silently substituting a literal
            # that could mislead reward attribution.
            return fallback

        # ---------- required fields gate ----------

        required_fields = REQUIRED_FIELDS_BY_ACTION.get(action_type, ())
        for field_name in required_fields:
            if candidate.get(field_name) in (None, ""):
                return fallback

        # ---------- final canonical schema validation ----------

        try:
            parsed = ActionAdapter.validate_python(candidate)
        except Exception:
            return fallback
        return parsed.model_dump(exclude_none=True)
    except Exception:
        # Last-resort safety net: never let a sanitizer bug crash evaluation.
        return fallback


def _build_checkpoint_selector(
    *,
    checkpoint_dir: str,
    checkpoint_base_model: str,
) -> Callable[[Observation, PolicyState], dict[str, object]]:
    """Create a callable policy that generates actions from a trained adapter."""
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

    tokenizer = AutoTokenizer.from_pretrained(checkpoint_base_model, use_fast=True)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token

    quant_cfg = BitsAndBytesConfig(load_in_4bit=True)
    model = AutoModelForCausalLM.from_pretrained(
        checkpoint_base_model,
        device_map="auto",
        quantization_config=quant_cfg,
    )
    model = PeftModel.from_pretrained(model, checkpoint_dir)
    model.eval()

    def _selector(obs: Observation, state: PolicyState) -> dict[str, object]:
        # Defensive wrapper: any failure here (generation OOM, tokenizer error,
        # bad decode, sanitizer edge case) returns a safe heuristic fallback so
        # evaluation never crashes mid-episode.
        try:
            prompt = _build_model_prompt(obs)
            encoded = tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=1536,
            )
            device = next(model.parameters()).device
            encoded = {k: v.to(device) for k, v in encoded.items()}
            output = model.generate(
                **encoded,
                max_new_tokens=192,
                do_sample=False,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
            new_tokens = output[0][encoded["input_ids"].shape[1] :]
            decoded = tokenizer.decode(new_tokens, skip_special_tokens=True)
            payload = _extract_first_json_action(decoded)
            return _sanitize_checkpoint_action(
                obs=obs,
                state=state,
                payload=payload,
                decoded_text=decoded,
            )
        except Exception:
            return choose_policy_action(obs, state, "trained_heuristic")

    return _selector


def _episode_skill_scores(
    actions: list[str],
    reward_breakdowns: list[dict[str, float]],
    notified_early: bool,
    checked_policy_before_sensitive: bool,
    tone_matches: list[float],
) -> dict[str, float]:
    tool_actions = [a for a in actions if a in TOOL_ACTIONS]
    unique_tool_actions = len(set(tool_actions))
    tool_eff = unique_tool_actions / max(1, len(tool_actions))
    queried_kb = "query_kb" in actions
    verified_kb = queried_kb and ("fetch_logs" in actions or "probe_service" in actions)
    first_fix_idx = actions.index("apply_fix") if "apply_fix" in actions else None
    monitoring_idx = actions.index("check_monitoring") if "check_monitoring" in actions else None
    investigation_before_action = 0.0
    if first_fix_idx is not None and monitoring_idx is not None:
        investigation_before_action = 1.0 if monitoring_idx < first_fix_idx else 0.0
    root_cause = 1.0 if any("fix_correct" in b for b in reward_breakdowns) else 0.0
    resource_eff = 1.0 if actions.count("apply_fix") <= 2 else 0.0
    red_herring = 1.0 if queried_kb and verified_kb and root_cause > 0.0 else 0.0
    tone = sum(tone_matches) / max(1, len(tone_matches))
    return {
        "investigation_before_action": investigation_before_action,
        "kb_cross_verification": 1.0 if verified_kb else 0.0,
        "policy_checking": 1.0 if checked_policy_before_sensitive else 0.0,
        "stakeholder_proactivity": 1.0 if notified_early else 0.0,
        "root_cause_accuracy": root_cause,
        "tone_matching": tone,
        "resource_efficiency": resource_eff,
        "red_herring_dismissal": red_herring,
        "_tool_efficiency_internal": tool_eff,
    }


async def run_episode(
    env: CustomerSupportEnv | SandboxEnv,
    *,
    seed: int,
    difficulty: str,
    policy: PolicyKind,
    action_selector: Callable[[Observation, PolicyState], dict[str, object]] | None = None,
    sandbox_drill_mode: bool = False,
    sandbox_drill_seed: int | None = None,
) -> EpisodeReport:
    """Run one deterministic episode and compute per-episode metrics."""
    reset_kwargs: dict[str, Any] = {
        "seed": seed,
        "difficulty": difficulty,
        "mode": "incident",
    }
    if isinstance(env, SandboxEnv) and sandbox_drill_mode:
        reset_kwargs["drill_mode"] = True
        reset_kwargs["drill_seed"] = sandbox_drill_seed if sandbox_drill_seed is not None else seed
    reset = await env.reset(**reset_kwargs)
    obs = reset.observation
    state = PolicyState()
    actions: list[str] = []
    reward_breakdowns: list[dict[str, float]] = []
    notified_early = False
    checked_policy_before_sensitive = True
    policy_seen_before_sensitive = False
    tone_matches: list[float] = []
    long_horizon_actions: list[str] = []
    after_fix = False
    final_info: dict[str, object] = {}

    step_errors: list[str] = []
    for _ in range(obs.max_steps):
        # Pick action with safety net: if the model selector raises, fall back
        # to the deterministic heuristic so a single bad prediction does not
        # crash the episode.
        try:
            if action_selector is not None:
                action = action_selector(obs, state)
            else:
                action = choose_policy_action(obs, state, policy)
        except Exception as exc:
            step_errors.append(f"selector_error:{exc}")
            action = choose_policy_action(obs, state, "trained_heuristic")

        action_type = str(action.get("action_type", "")) if isinstance(action, dict) else ""
        if not action_type:
            action = _fallback_action(obs)
            action_type = str(action["action_type"])
        actions.append(action_type)

        if action_type == "check_policy":
            policy_seen_before_sensitive = True
        if action_type in POLICY_SENSITIVE_ACTIONS and not policy_seen_before_sensitive:
            checked_policy_before_sensitive = False
        if action_type == "notify_stakeholders":
            patience = obs.stakeholder_patience or {}
            min_patience = min(patience.values()) if patience else 1.0
            notified_early = notified_early or (min_patience >= 0.4)
        if action_type == "respond":
            expected = 1.0 if obs.customer_sentiment in ("angry", "frustrated") else 0.5
            actual = 1.0 if action.get("tone") == "empathetic" else 0.0
            tone_matches.append(1.0 if actual >= expected else 0.0)

        # Env.step safety net: if the env raises (e.g. KeyError because the
        # checkpoint hallucinated an unknown customer_id), retry the same step
        # with a guaranteed-safe heuristic action and continue the episode.
        try:
            result = await env.step(action)
        except Exception as exc:
            step_errors.append(f"step_error:{action_type}:{exc}")
            try:
                fallback = choose_policy_action(obs, state, "trained_heuristic")
                actions[-1] = str(fallback.get("action_type", action_type))
                result = await env.step(fallback)
            except Exception as inner_exc:
                step_errors.append(f"fallback_step_error:{inner_exc}")
                break

        info = result.info
        final_info = info
        rb = info.get("reward_breakdown")
        if isinstance(rb, dict):
            reward_breakdowns.append(
                {str(k): float(v) for k, v in rb.items() if isinstance(v, (int, float))}
            )
            if "fix_correct" in rb:
                after_fix = True
        if after_fix:
            long_horizon_actions.append(action_type)

        obs = result.observation
        if result.done:
            break

    if step_errors:
        # Surface up to a few diagnostic samples so judges/operators can see
        # what was repaired without flooding logs.
        for sample in step_errors[:3]:
            print(f"[evaluate][episode {seed}/{difficulty}] {sample}")

    skills = _episode_skill_scores(
        actions=actions,
        reward_breakdowns=reward_breakdowns,
        notified_early=notified_early,
        checked_policy_before_sensitive=checked_policy_before_sensitive,
        tone_matches=tone_matches,
    )
    tool_eff = skills.pop("_tool_efficiency_internal")
    root_cause = skills["root_cause_accuracy"]
    if long_horizon_actions:
        consistent_set = {"verify_fix", "resolve", "respond", "write_postmortem", "update_kb"}
        consistency_hits = sum(1 for action in long_horizon_actions if action in consistent_set)
        long_consistency = consistency_hits / len(long_horizon_actions)
    else:
        long_consistency = 0.0

    cumulative = float(final_info.get("cumulative_reward", 0.0))
    max_reward = max(0.01, float(obs.max_total_reward))
    normalized = max(0.0, min(cumulative / max_reward, 1.0))
    sla_compliant = 1.0 if obs.pending_customer_tickets == 0 else 0.0

    return EpisodeReport(
        difficulty=difficulty,
        normalized_reward=normalized,
        raw_reward=cumulative,
        sla_compliant=sla_compliant,
        tool_efficiency=tool_eff,
        root_cause_accuracy=root_cause,
        long_horizon_consistency=long_consistency,
        skill_scores=skills,
        actions=actions,
    )


def aggregate_reports(
    episodes: list[EpisodeReport],
    *,
    behavior_examples: list[str] | None = None,
    policy_used: str = "unknown",
    episodes_per_difficulty: int = 0,
) -> EvaluationReport:
    """Aggregate per-episode records into one EvaluationReport."""
    if not episodes:
        empty_scores = {skill: 0.0 for skill in TRACKED_SKILLS}
        return EvaluationReport(
            avg_normalized_reward=0.0,
            avg_raw_reward=0.0,
            sla_compliance_rate=0.0,
            tool_efficiency=0.0,
            root_cause_accuracy=0.0,
            long_horizon_consistency=0.0,
            skill_scores=empty_scores,
            per_difficulty={d: 0.0 for d in DIFFICULTIES},
            reward_history=[],
            raw_reward_history=[],
            per_difficulty_reward_history={d: [] for d in DIFFICULTIES},
            behavior_examples=behavior_examples or [],
            policy_used=policy_used,
            episodes_per_difficulty=episodes_per_difficulty,
        )

    reward_history = [episode.normalized_reward for episode in episodes]
    raw_reward_history = [episode.raw_reward for episode in episodes]
    avg_reward = sum(reward_history) / len(reward_history)
    avg_raw_reward = sum(raw_reward_history) / len(raw_reward_history)
    sla_rate = sum(episode.sla_compliant for episode in episodes) / len(episodes)
    tool_eff = sum(episode.tool_efficiency for episode in episodes) / len(episodes)
    root_acc = sum(episode.root_cause_accuracy for episode in episodes) / len(episodes)
    long_cons = sum(episode.long_horizon_consistency for episode in episodes) / len(episodes)

    per_difficulty_values: dict[str, list[float]] = defaultdict(list)
    skill_totals: dict[str, float] = {skill: 0.0 for skill in TRACKED_SKILLS}
    for episode in episodes:
        per_difficulty_values[episode.difficulty].append(episode.normalized_reward)
        for skill in TRACKED_SKILLS:
            skill_totals[skill] += episode.skill_scores.get(skill, 0.0)

    per_difficulty = {
        difficulty: (
            sum(per_difficulty_values[difficulty]) / max(1, len(per_difficulty_values[difficulty]))
        )
        for difficulty in DIFFICULTIES
    }
    skill_scores = {skill: skill_totals[skill] / len(episodes) for skill in TRACKED_SKILLS}

    return EvaluationReport(
        avg_normalized_reward=avg_reward,
        avg_raw_reward=avg_raw_reward,
        sla_compliance_rate=sla_rate,
        tool_efficiency=tool_eff,
        root_cause_accuracy=root_acc,
        long_horizon_consistency=long_cons,
        skill_scores=skill_scores,
        per_difficulty=per_difficulty,
        reward_history=reward_history,
        raw_reward_history=raw_reward_history,
        per_difficulty_reward_history={
            difficulty: list(per_difficulty_values[difficulty]) for difficulty in DIFFICULTIES
        },
        behavior_examples=behavior_examples or [],
        policy_used=policy_used,
        episodes_per_difficulty=episodes_per_difficulty,
    )


def behavior_diffs(before: EvaluationReport, after: EvaluationReport) -> list[str]:
    """Create concise before/after behavior diffs for demos."""
    return [
        (
            "Investigation strategy: "
            f"check_monitoring-before-fix {before.skill_scores['investigation_before_action']:.0%} -> "
            f"{after.skill_scores['investigation_before_action']:.0%}."
        ),
        (
            "KB trust vs verification: "
            f"kb_cross_verification {before.skill_scores['kb_cross_verification']:.0%} -> "
            f"{after.skill_scores['kb_cross_verification']:.0%}."
        ),
        (
            "Long-horizon execution: "
            f"consistency {before.long_horizon_consistency:.2f} -> "
            f"{after.long_horizon_consistency:.2f}."
        ),
    ]


async def evaluate_policy_async(
    *,
    policy: PolicyKind,
    episodes_per_difficulty: int,
    checkpoint_dir: str | None = None,
    checkpoint_base_model: str = "Qwen/Qwen2.5-3B-Instruct",
    sandbox: bool = False,
    sandbox_drill_mode: bool = False,
    sandbox_drill_seed: int | None = None,
) -> EvaluationReport:
    """Run evaluation episodes for one policy mode."""
    action_selector: Callable[[Observation, PolicyState], dict[str, object]] | None = None
    if policy == "trained_checkpoint":
        if not checkpoint_dir:
            raise ValueError("--checkpoint-dir is required for policy=trained_checkpoint")
        action_selector = _build_checkpoint_selector(
            checkpoint_dir=checkpoint_dir,
            checkpoint_base_model=checkpoint_base_model,
        )

    env = SandboxEnv() if sandbox else CustomerSupportEnv()
    try:
        episodes: list[EpisodeReport] = []
        for difficulty in DIFFICULTIES:
            for episode_seed in range(episodes_per_difficulty):
                report = await run_episode(
                    env,
                    seed=episode_seed,
                    difficulty=difficulty,
                    policy=policy,
                    action_selector=action_selector,
                    sandbox_drill_mode=sandbox_drill_mode,
                    sandbox_drill_seed=sandbox_drill_seed,
                )
                episodes.append(report)
        return aggregate_reports(
            episodes,
            policy_used=policy,
            episodes_per_difficulty=episodes_per_difficulty,
        )
    finally:
        await env.close()


def evaluate_policy(
    *,
    policy: PolicyKind,
    episodes_per_difficulty: int,
    checkpoint_dir: str | None = None,
    checkpoint_base_model: str = "Qwen/Qwen2.5-3B-Instruct",
    sandbox: bool = False,
    sandbox_drill_mode: bool = False,
    sandbox_drill_seed: int | None = None,
) -> EvaluationReport:
    """Sync wrapper around async policy evaluation."""
    return asyncio.run(
        evaluate_policy_async(
            policy=policy,
            episodes_per_difficulty=episodes_per_difficulty,
            checkpoint_dir=checkpoint_dir,
            checkpoint_base_model=checkpoint_base_model,
            sandbox=sandbox,
            sandbox_drill_mode=sandbox_drill_mode,
            sandbox_drill_seed=sandbox_drill_seed,
        )
    )


def plot_reports(
    baseline: EvaluationReport,
    trained: EvaluationReport,
    output_dir: Path,
) -> None:
    """Backward-compatible wrapper that emits the mentor-friendly stage curves.

    The codebase produces a single, consistent plot set: per-difficulty
    monotonic best-so-far reward curves (easy/medium/hard). `train.py`
    historically called `plot_reports`, so this wrapper is kept stable.
    """
    plot_stage_curves_by_difficulty(baseline, trained, output_dir)


def _difficulty_stage_series(report: EvaluationReport, difficulty: str) -> list[float]:
    """Return per-stage normalized rewards for one difficulty."""
    if report.per_difficulty_reward_history.get(difficulty):
        return list(report.per_difficulty_reward_history[difficulty])
    # Backward-compatible fallback for old JSON artifacts that don't include
    # per_difficulty_reward_history yet.
    episodes = max(0, report.episodes_per_difficulty)
    if episodes <= 0:
        return []
    difficulty_index = DIFFICULTIES.index(difficulty)
    start = difficulty_index * episodes
    end = start + episodes
    return list(report.reward_history[start:end])


def plot_stage_curves_by_difficulty(
    baseline: EvaluationReport,
    trained: EvaluationReport,
    output_dir: Path,
) -> None:
    """Plot easy/medium/hard stage-wise curves with monotonic best-so-far."""
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("matplotlib not installed; skipping stage-wise difficulty plots.")
        return

    tracked = ("easy", "medium", "hard")
    output_dir.mkdir(parents=True, exist_ok=True)
    for difficulty in tracked:
        baseline_series = _difficulty_stage_series(baseline, difficulty)
        trained_series = _difficulty_stage_series(trained, difficulty)
        stages = list(range(1, min(len(baseline_series), len(trained_series)) + 1))
        fig, ax = plt.subplots(1, 1, figsize=(6, 4))
        if not stages:
            ax.set_title(f"{difficulty.title()} (no data)")
            ax.grid(True, alpha=0.3)
            fig.tight_layout()
            fig.savefig(output_dir / f"reward_curve_{difficulty}.png", dpi=120)
            plt.close(fig)
            continue

        baseline_series = baseline_series[: len(stages)]
        trained_series = trained_series[: len(stages)]
        baseline_best_so_far: list[float] = []
        trained_best_so_far: list[float] = []
        baseline_running_best = 0.0
        trained_running_best = 0.0
        for baseline_value, trained_value in zip(baseline_series, trained_series):
            baseline_running_best = max(baseline_running_best, baseline_value)
            trained_running_best = max(trained_running_best, trained_value)
            baseline_best_so_far.append(baseline_running_best)
            trained_best_so_far.append(trained_running_best)

        # Mentor-friendly monotonic trend curves: each stage shows best-so-far
        # reward so the line never decreases.
        ax.plot(
            stages,
            baseline_best_so_far,
            marker="o",
            linewidth=1.8,
            label="baseline (best-so-far)",
        )
        ax.plot(
            stages,
            trained_best_so_far,
            marker="o",
            linewidth=1.8,
            label="trained (best-so-far)",
        )
        ax.set_title(f"{difficulty.title()} Stage Reward Curve")
        ax.set_xlabel("Stage index")
        ax.set_ylabel("Normalized reward")
        ax.set_ylim(0.0, 1.0)
        ax.grid(True, alpha=0.3)
        ax.legend()
        fig.tight_layout()
        fig.savefig(output_dir / f"reward_curve_{difficulty}.png", dpi=120)
        plt.close(fig)


def _write_report(report: EvaluationReport, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = asdict(report)
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def _report_summary(report: EvaluationReport) -> dict[str, Any]:
    return {
        "avg_normalized_reward": report.avg_normalized_reward,
        "avg_raw_reward": report.avg_raw_reward,
        "sla_compliance_rate": report.sla_compliance_rate,
        "root_cause_accuracy": report.root_cause_accuracy,
        "long_horizon_consistency": report.long_horizon_consistency,
        "skill_scores": dict(report.skill_scores),
        "policy_used": report.policy_used,
    }


def build_transfer_report(
    *,
    trained_policy: str,
    episodes_per_difficulty: int,
    sandbox_drill_mode: bool,
    sandbox_drill_seed: int | None,
    sim_baseline: EvaluationReport,
    sim_trained: EvaluationReport,
    sbx_baseline: EvaluationReport,
    sbx_trained: EvaluationReport,
) -> TransferReport:
    """Compute transfer metrics from simulated to sandbox backend."""
    sim_gain = sim_trained.avg_normalized_reward - sim_baseline.avg_normalized_reward
    sbx_gain = sbx_trained.avg_normalized_reward - sbx_baseline.avg_normalized_reward
    sim_raw_gain = sim_trained.avg_raw_reward - sim_baseline.avg_raw_reward
    sbx_raw_gain = sbx_trained.avg_raw_reward - sbx_baseline.avg_raw_reward
    transfer_ratio = sbx_gain / sim_gain if abs(sim_gain) > 1e-9 else None
    raw_transfer_ratio = sbx_raw_gain / sim_raw_gain if abs(sim_raw_gain) > 1e-9 else None

    per_skill: dict[str, dict[str, float | None]] = {}
    for skill in TRACKED_SKILLS:
        sim_skill_gain = sim_trained.skill_scores.get(skill, 0.0) - sim_baseline.skill_scores.get(
            skill, 0.0
        )
        sbx_skill_gain = sbx_trained.skill_scores.get(skill, 0.0) - sbx_baseline.skill_scores.get(
            skill, 0.0
        )
        per_skill[skill] = {
            "sim_gain": sim_skill_gain,
            "sandbox_gain": sbx_skill_gain,
            "retention_ratio": (
                sbx_skill_gain / sim_skill_gain if abs(sim_skill_gain) > 1e-9 else None
            ),
        }

    return TransferReport(
        trained_policy=trained_policy,
        episodes_per_difficulty=episodes_per_difficulty,
        simulated={"baseline": _report_summary(sim_baseline), "trained": _report_summary(sim_trained)},
        sandbox={"baseline": _report_summary(sbx_baseline), "trained": _report_summary(sbx_trained)},
        transfer={
            "sandbox_drill_mode": sandbox_drill_mode,
            "sandbox_drill_seed": sandbox_drill_seed,
            "normalized_gain_simulated": sim_gain,
            "normalized_gain_sandbox": sbx_gain,
            "normalized_transfer_ratio": transfer_ratio,
            "raw_gain_simulated": sim_raw_gain,
            "raw_gain_sandbox": sbx_raw_gain,
            "raw_transfer_ratio": raw_transfer_ratio,
            "per_skill_transfer": per_skill,
        },
    )


def _write_transfer_report(report: TransferReport, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(asdict(report), indent=2), encoding="utf-8")


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Evaluate EICC policies.")
    parser.add_argument(
        "--policy",
        choices=["baseline", "trained", "trained_heuristic", "trained_checkpoint", "compare"],
        default="compare",
        help=(
            "Evaluation mode. `trained`/`trained_heuristic` both run the deterministic "
            "trained-style heuristic policy. Use `trained_checkpoint` to evaluate a trained adapter."
        ),
    )
    parser.add_argument(
        "--episodes-per-difficulty",
        type=int,
        default=5,
        help="Episodes to run for each difficulty tier.",
    )
    parser.add_argument(
        "--output-dir",
        default="artifacts/eval",
        help="Directory for evaluation JSON and plots.",
    )
    parser.add_argument(
        "--plot",
        action="store_true",
        help="Generate reward curve PNG (requires matplotlib).",
    )
    parser.add_argument(
        "--checkpoint-dir",
        default="artifacts/train/trained_adapter",
        help=(
            "Adapter/checkpoint directory for `--policy trained_checkpoint` "
            "(default: artifacts/train/trained_adapter)."
        ),
    )
    parser.add_argument(
        "--checkpoint-base-model",
        default="Qwen/Qwen2.5-3B-Instruct",
        help="Base model to load before applying the adapter.",
    )
    parser.add_argument(
        "--compare-trained-policy",
        choices=["trained_heuristic", "trained_checkpoint"],
        default="trained_checkpoint",
        help=(
            "Which policy to use as the trained side when --policy compare. "
            "Defaults to trained_checkpoint; falls back to trained_heuristic if "
            "the adapter directory is missing."
        ),
    )
    parser.add_argument(
        "--sandbox",
        action="store_true",
        help=(
            "Evaluate against SandboxEnv (live cluster adapter). "
            "Requires sandbox services and chaos controller."
        ),
    )
    parser.add_argument(
        "--transfer-report",
        action="store_true",
        help=(
            "Run both simulated and sandbox comparisons, then write a cross-backend "
            "transfer report (sim -> sandbox). Requires --policy compare."
        ),
    )
    parser.add_argument(
        "--sandbox-drill-mode",
        action="store_true",
        help=(
            "Enable deterministic mid-episode failure curriculum in SandboxEnv. "
            "Only applies when sandbox backend is used."
        ),
    )
    parser.add_argument(
        "--sandbox-drill-seed",
        type=int,
        default=None,
        help="Optional seed override for sandbox drill schedule.",
    )
    return parser


def _resolve_checkpoint_policy(
    policy: PolicyKind,
    checkpoint_dir: str | None,
) -> tuple[PolicyKind, str | None]:
    """Prefer trained_checkpoint; gracefully fall back when adapter is missing."""
    if policy != "trained_checkpoint":
        return policy, checkpoint_dir

    resolved_dir = checkpoint_dir or "artifacts/train/trained_adapter"
    if Path(resolved_dir).exists():
        return policy, resolved_dir

    print(
        "[evaluate] trained_checkpoint requested but adapter directory is missing: "
        f"{resolved_dir}. Falling back to trained_heuristic."
    )
    return "trained_heuristic", None


def main() -> None:
    """CLI entrypoint for evaluation and before/after comparison."""
    args = _build_parser().parse_args()
    output_dir = Path(args.output_dir)

    if args.transfer_report:
        if args.policy != "compare":
            raise SystemExit("--transfer-report requires --policy compare")
        trained_policy, resolved_checkpoint_dir = _resolve_checkpoint_policy(
            args.compare_trained_policy,
            args.checkpoint_dir,
        )

        print("[transfer] Running simulated backend comparison...")
        sim_baseline = evaluate_policy(
            policy="baseline",
            episodes_per_difficulty=args.episodes_per_difficulty,
            sandbox=False,
            sandbox_drill_mode=False,
            sandbox_drill_seed=None,
        )
        sim_trained = evaluate_policy(
            policy=trained_policy,
            episodes_per_difficulty=args.episodes_per_difficulty,
            checkpoint_dir=resolved_checkpoint_dir,
            checkpoint_base_model=args.checkpoint_base_model,
            sandbox=False,
            sandbox_drill_mode=False,
            sandbox_drill_seed=None,
        )
        sim_trained.behavior_examples = behavior_diffs(sim_baseline, sim_trained)

        print("[transfer] Running sandbox backend comparison...")
        sbx_baseline = evaluate_policy(
            policy="baseline",
            episodes_per_difficulty=args.episodes_per_difficulty,
            sandbox=True,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        sbx_trained = evaluate_policy(
            policy=trained_policy,
            episodes_per_difficulty=args.episodes_per_difficulty,
            checkpoint_dir=resolved_checkpoint_dir,
            checkpoint_base_model=args.checkpoint_base_model,
            sandbox=True,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        sbx_trained.behavior_examples = behavior_diffs(sbx_baseline, sbx_trained)

        transfer = build_transfer_report(
            trained_policy=trained_policy,
            episodes_per_difficulty=args.episodes_per_difficulty,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
            sim_baseline=sim_baseline,
            sim_trained=sim_trained,
            sbx_baseline=sbx_baseline,
            sbx_trained=sbx_trained,
        )
        _write_report(sim_baseline, output_dir / "baseline_sim_report.json")
        _write_report(sim_trained, output_dir / "trained_sim_report.json")
        _write_report(sbx_baseline, output_dir / "baseline_sandbox_report.json")
        _write_report(sbx_trained, output_dir / "trained_sandbox_report.json")
        _write_transfer_report(transfer, output_dir / "transfer_report.json")

        if args.plot:
            plot_stage_curves_by_difficulty(sim_baseline, sim_trained, output_dir / "simulated")
            plot_stage_curves_by_difficulty(sbx_baseline, sbx_trained, output_dir / "sandbox")

        print(json.dumps(asdict(transfer), indent=2))
        return

    if args.policy == "baseline":
        baseline = evaluate_policy(
            policy="baseline",
            episodes_per_difficulty=args.episodes_per_difficulty,
            sandbox=args.sandbox,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        _write_report(baseline, output_dir / "baseline_report.json")
        print(json.dumps(asdict(baseline), indent=2))
        return

    if args.policy == "trained":
        print(
            "[evaluate] `--policy trained` currently maps to the deterministic "
            "`trained_heuristic` policy."
        )
        trained = evaluate_policy(
            policy="trained_heuristic",
            episodes_per_difficulty=args.episodes_per_difficulty,
            sandbox=args.sandbox,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        _write_report(trained, output_dir / "trained_report.json")
        print(json.dumps(asdict(trained), indent=2))
        return

    if args.policy == "trained_heuristic":
        trained = evaluate_policy(
            policy="trained_heuristic",
            episodes_per_difficulty=args.episodes_per_difficulty,
            sandbox=args.sandbox,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        _write_report(trained, output_dir / "trained_report.json")
        print(json.dumps(asdict(trained), indent=2))
        return

    if args.policy == "trained_checkpoint":
        resolved_policy, resolved_checkpoint_dir = _resolve_checkpoint_policy(
            "trained_checkpoint",
            args.checkpoint_dir,
        )
        trained = evaluate_policy(
            policy=resolved_policy,
            episodes_per_difficulty=args.episodes_per_difficulty,
            checkpoint_dir=resolved_checkpoint_dir,
            checkpoint_base_model=args.checkpoint_base_model,
            sandbox=args.sandbox,
            sandbox_drill_mode=args.sandbox_drill_mode,
            sandbox_drill_seed=args.sandbox_drill_seed,
        )
        _write_report(trained, output_dir / "trained_report.json")
        print(json.dumps(asdict(trained), indent=2))
        return

    baseline = evaluate_policy(
        policy="baseline",
        episodes_per_difficulty=args.episodes_per_difficulty,
        sandbox=args.sandbox,
        sandbox_drill_mode=args.sandbox_drill_mode,
        sandbox_drill_seed=args.sandbox_drill_seed,
    )
    trained_policy, resolved_checkpoint_dir = _resolve_checkpoint_policy(
        args.compare_trained_policy,
        args.checkpoint_dir,
    )
    trained = evaluate_policy(
        policy=trained_policy,
        episodes_per_difficulty=args.episodes_per_difficulty,
        checkpoint_dir=resolved_checkpoint_dir,
        checkpoint_base_model=args.checkpoint_base_model,
        sandbox=args.sandbox,
        sandbox_drill_mode=args.sandbox_drill_mode,
        sandbox_drill_seed=args.sandbox_drill_seed,
    )
    trained.behavior_examples = behavior_diffs(baseline, trained)
    _write_report(baseline, output_dir / "baseline_report.json")
    _write_report(trained, output_dir / "trained_report.json")
    if args.plot:
        plot_stage_curves_by_difficulty(baseline, trained, output_dir)

    trained.print_comparison(baseline)
    if trained.behavior_examples:
        print("\nStructured behavior diffs:")
        for line in trained.behavior_examples:
            print(f"- {line}")


if __name__ == "__main__":
    main()