| """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_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 |
|
|
|
|
| |
| _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 {} |
| |
| 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) |
|
|
| |
| 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"} |
| |
| 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 |
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
|
|
| if action_type == "follow_runbook_step": |
| |
| if candidate.get("step_index") in (None, "") and candidate.get("step") not in (None, ""): |
| try: |
| candidate["step_index"] = int(candidate["step"]) |
| except (TypeError, ValueError): |
| pass |
| |
| |
| 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 |
| |
| 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")) |
| 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"}: |
| |
| 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(): |
| |
| 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_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 |
| |
| |
| |
| return fallback |
|
|
| |
|
|
| required_fields = REQUIRED_FIELDS_BY_ACTION.get(action_type, ()) |
| for field_name in required_fields: |
| if candidate.get(field_name) in (None, ""): |
| return fallback |
|
|
| |
|
|
| try: |
| parsed = ActionAdapter.validate_python(candidate) |
| except Exception: |
| return fallback |
| return parsed.model_dump(exclude_none=True) |
| except Exception: |
| |
| 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]: |
| |
| |
| |
| 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): |
| |
| |
| |
| 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) |
|
|
| |
| |
| |
| 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: |
| |
| |
| 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]) |
| |
| |
| 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) |
|
|
| |
| |
| 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() |
|
|
