policy_learning.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import importlib
import json
from dataclasses import asdict, dataclass, field
from pathlib import Path
from statistics import mean
from typing import Any, Callable, Iterable

from models import HelpdeskTicketAction, HelpdeskTicketObservation
from server.environment import HelpdeskTicketRoutingEnvironment
from server.tasks import get_task_definition
from vocabulary import TASK_IDS


DEFAULT_COMPARE_POLICIES = (
    "no_investigation",
    "investigate_when_context_hidden",
    "adaptive_cue_bandit",
)
DEFAULT_SEARCH_POLICIES = (
    "no_investigation",
    "legacy_single_probe",
    "investigate_when_context_hidden",
    "adaptive_cue_bandit",
)
DEFAULT_OUTPUT_DIR = "analysis/policy_learning_runs"

SubmitBuilder = Callable[[dict[str, Any], list[str]], HelpdeskTicketAction]
EnvFactory = Callable[[], HelpdeskTicketRoutingEnvironment]


@dataclass(frozen=True)
class PolicyConfig:
    name: str
    investigate_hidden_context: bool
    investigate_related_ticket_hint: bool
    investigate_ambiguity_history: bool
    max_investigations_per_ticket: int
    description: str
    strategy: str = "static"


POLICY_LIBRARY: dict[str, PolicyConfig] = {
    "no_investigation": PolicyConfig(
        name="no_investigation",
        strategy="static",
        investigate_hidden_context=False,
        investigate_related_ticket_hint=False,
        investigate_ambiguity_history=False,
        max_investigations_per_ticket=0,
        description="Always submit immediately and never investigate.",
    ),
    "legacy_single_probe": PolicyConfig(
        name="legacy_single_probe",
        strategy="static",
        investigate_hidden_context=False,
        investigate_related_ticket_hint=True,
        investigate_ambiguity_history=True,
        max_investigations_per_ticket=1,
        description="Mimics the earlier single-tool hint policy.",
    ),
    "investigate_when_context_hidden": PolicyConfig(
        name="investigate_when_context_hidden",
        strategy="static",
        investigate_hidden_context=True,
        investigate_related_ticket_hint=False,
        investigate_ambiguity_history=False,
        max_investigations_per_ticket=1,
        description="Investigate once when the environment shows hidden-context pressure.",
    ),
    "adaptive_cue_bandit": PolicyConfig(
        name="adaptive_cue_bandit",
        strategy="adaptive",
        investigate_hidden_context=True,
        investigate_related_ticket_hint=True,
        investigate_ambiguity_history=True,
        max_investigations_per_ticket=3,
        description=(
            "Learn cue-conditioned tool preferences from investigation rewards on train seeds."
        ),
    ),
}

AVAILABLE_TOOLS = (
    "lookup_related_ticket",
    "lookup_requester_history",
    "lookup_internal_routing_note",
    "lookup_queue_capacity_forecast",
    "lookup_queue_cluster_summary",
)


@dataclass
class AdaptiveToolBandit:
    exploration_rounds: int = 1
    cue_tool_totals: dict[str, dict[str, float]] = field(default_factory=dict)
    cue_tool_counts: dict[str, dict[str, int]] = field(default_factory=dict)
    global_tool_totals: dict[str, float] = field(default_factory=dict)
    global_tool_counts: dict[str, int] = field(default_factory=dict)

    def choose_tool(self, cue: str, candidate_tools: list[str]) -> str:
        for tool_name in candidate_tools:
            if self.cue_tool_counts.get(cue, {}).get(tool_name, 0) < self.exploration_rounds:
                return tool_name
        return max(
            candidate_tools,
            key=lambda tool_name: (
                self._cue_average(cue, tool_name),
                self._global_average(tool_name),
                -candidate_tools.index(tool_name),
            ),
        )

    def record_reward(self, cue: str, tool_name: str, reward: float) -> None:
        cue_totals = self.cue_tool_totals.setdefault(cue, {})
        cue_counts = self.cue_tool_counts.setdefault(cue, {})
        cue_totals[tool_name] = cue_totals.get(tool_name, 0.0) + reward
        cue_counts[tool_name] = cue_counts.get(tool_name, 0) + 1
        self.global_tool_totals[tool_name] = self.global_tool_totals.get(tool_name, 0.0) + reward
        self.global_tool_counts[tool_name] = self.global_tool_counts.get(tool_name, 0) + 1

    def export(self) -> dict[str, Any]:
        return {
            "exploration_rounds": self.exploration_rounds,
            "cue_tool_averages": {
                cue: {
                    tool_name: round(self._cue_average(cue, tool_name), 6)
                    for tool_name in sorted(tool_totals)
                }
                for cue, tool_totals in sorted(self.cue_tool_totals.items())
            },
            "global_tool_averages": {
                tool_name: round(self._global_average(tool_name), 6)
                for tool_name in sorted(self.global_tool_totals)
            },
        }

    def frozen_copy(self) -> "AdaptiveToolBandit":
        return AdaptiveToolBandit(
            exploration_rounds=self.exploration_rounds,
            cue_tool_totals={
                cue: dict(tool_totals) for cue, tool_totals in self.cue_tool_totals.items()
            },
            cue_tool_counts={
                cue: dict(tool_counts) for cue, tool_counts in self.cue_tool_counts.items()
            },
            global_tool_totals=dict(self.global_tool_totals),
            global_tool_counts=dict(self.global_tool_counts),
        )

    def _cue_average(self, cue: str, tool_name: str) -> float:
        total = self.cue_tool_totals.get(cue, {}).get(tool_name, 0.0)
        count = self.cue_tool_counts.get(cue, {}).get(tool_name, 0)
        if count == 0:
            return self._global_average(tool_name)
        return total / count

    def _global_average(self, tool_name: str) -> float:
        total = self.global_tool_totals.get(tool_name, 0.0)
        count = self.global_tool_counts.get(tool_name, 0)
        if count == 0:
            return 0.0
        return total / count


def _dedupe_preserving_order(values: Iterable[int]) -> list[int]:
    seen: set[int] = set()
    ordered: list[int] = []
    for value in values:
        if value in seen:
            continue
        seen.add(value)
        ordered.append(value)
    return ordered


def parse_int_spec(spec: str, *, field_name: str) -> list[int]:
    values: list[int] = []
    for chunk in spec.split(","):
        part = chunk.strip()
        if not part:
            continue
        if "-" in part:
            start_raw, end_raw = part.split("-", 1)
            try:
                start = int(start_raw)
                end = int(end_raw)
            except ValueError as exc:
                raise ValueError(f"{field_name} contains an invalid range: {part!r}") from exc
            if end < start:
                raise ValueError(f"{field_name} range must be ascending: {part!r}")
            values.extend(range(start, end + 1))
            continue
        try:
            values.append(int(part))
        except ValueError as exc:
            raise ValueError(f"{field_name} contains an invalid integer: {part!r}") from exc
    if not values:
        raise ValueError(f"{field_name} must not be empty")
    return _dedupe_preserving_order(values)


def parse_task_ids(spec: str) -> list[int]:
    task_ids = parse_int_spec(spec, field_name="task_ids")
    unsupported = [task_id for task_id in task_ids if task_id not in TASK_IDS]
    if unsupported:
        raise ValueError(f"Unsupported task_ids: {unsupported}")
    return task_ids


def resolve_policies(spec: str) -> list[PolicyConfig]:
    names = [name.strip() for name in spec.split(",") if name.strip()]
    if not names:
        raise ValueError("At least one policy must be specified")
    policies: list[PolicyConfig] = []
    for name in names:
        if name not in POLICY_LIBRARY:
            raise ValueError(
                f"Unknown policy {name!r}. Available policies: {sorted(POLICY_LIBRARY)}"
            )
        policies.append(POLICY_LIBRARY[name])
    return policies


def default_submit_builder(
    ticket: dict[str, Any], allowed_fields: list[str]
) -> HelpdeskTicketAction:
    inference = importlib.import_module("inference")
    candidate = inference.heuristic_action(ticket, allowed_fields)
    candidate, _ = inference.apply_domain_overrides(ticket, candidate, allowed_fields)
    candidate, _ = inference.apply_capacity_planning_overrides(
        ticket,
        candidate,
        allowed_fields,
    )
    return HelpdeskTicketAction(**candidate)


def _routing_text(ticket: dict[str, Any]) -> str:
    parts = [
        str(ticket.get("title", "")),
        str(ticket.get("description", "")),
        str(ticket.get("ambiguity_note", "")),
        str(ticket.get("planning_note", "")),
        str(ticket.get("customer_update_note", "")),
        json.dumps(ticket.get("last_tool_result") or {}, sort_keys=True),
        json.dumps(ticket.get("routing_options") or [], sort_keys=True),
        json.dumps(ticket.get("operational_context") or {}, sort_keys=True),
        json.dumps(ticket.get("cluster_summary") or {}, sort_keys=True),
        json.dumps(ticket.get("capacity_state") or {}, sort_keys=True),
        json.dumps(ticket.get("future_queue_demand") or {}, sort_keys=True),
    ]
    related_preview = ticket.get("related_ticket_preview") or {}
    parts.extend(
        [
            str(related_preview.get("title", "")),
            str(related_preview.get("description", "")),
        ]
    )
    return " ".join(parts).lower()


def infer_ticket_cue(ticket: dict[str, Any]) -> str:
    text = _routing_text(ticket)
    context_status = ticket.get("context_status") or {}
    if (
        ticket.get("planning_note")
        or ticket.get("routing_options")
        or (ticket.get("operational_context") or {}).get("incident_recommended")
        or "lookup_queue_capacity_forecast"
        in (context_status.get("recommended_tools") or [])
        or any(
            phrase in text
            for phrase in (
                "capacity",
                "saturated",
                "backlog",
                "resource pressure",
                "alternate route",
            )
        )
    ):
        return "capacity_planning"
    if (
        bool((ticket.get("operational_context") or {}).get("cluster_coordination_hint"))
        or int((ticket.get("cluster_summary") or {}).get("future_cluster_ticket_count", 0) or 0)
        > 0
        or int((ticket.get("cluster_summary") or {}).get("shared_requester_count", 0) or 0)
        > 1
        or any(
            phrase in text
            for phrase in (
                "single coordinated owner",
                "existing workstream",
                "request cluster",
                "parallel workstream",
            )
        )
    ):
        return "cluster_coordination"
    if any(
        phrase in text
        for phrase in ("re:", "follow-up", "following up", "regression", "reference ticket", "third update")
    ):
        return "follow_up"
    if any(
        phrase in text
        for phrase in (
            "pricing",
            "quote",
            "vendor offer",
            "prorating",
            "seat expansion",
            "commercial",
        )
    ):
        return "commercial_ambiguity"
    if any(
        phrase in text
        for phrase in (
            "onboarding",
            "contractor",
            "permissions error",
            "blocked by an account problem",
        )
    ):
        return "workflow_blocker"
    if any(
        phrase in text
        for phrase in ("compliance scan", "vulnerability", "policy issue", "routing note")
    ):
        return "routing_note"
    if any(
        phrase in text
        for phrase in ("still", "again", "overdue", "legal", "priority")
    ):
        return "history_pressure"
    if any(phrase in text for phrase in ("incident", "outage", "lockout", "company-wide")):
        return "incident_pressure"
    return "generic_hidden_context"


def preferred_tool_order(
    ticket: dict[str, Any],
    *,
    hidden_context_remaining: bool,
) -> list[str]:
    text = _routing_text(ticket)
    context_status = ticket.get("context_status") or {}
    last_tool_result = ticket.get("last_tool_result") or {}
    last_tool_name = str(last_tool_result.get("tool_name", "") or "")
    recommended_tools = list(context_status.get("recommended_tools") or [])

    preferred_tools: list[str] = []
    if "lookup_queue_capacity_forecast" in recommended_tools:
        preferred_tools.append("lookup_queue_capacity_forecast")
    if last_tool_name == "lookup_related_ticket":
        preferred_tools.append("lookup_requester_history")
    if last_tool_name == "lookup_requester_history":
        preferred_tools.append("lookup_internal_routing_note")
    if last_tool_name == "lookup_internal_routing_note":
        preferred_tools.append("lookup_queue_cluster_summary")
    if last_tool_name == "lookup_queue_cluster_summary":
        preferred_tools.append("lookup_queue_capacity_forecast")

    if any(
        phrase in text
        for phrase in ("re:", "follow-up", "following up", "regression", "reference ticket")
    ) or ticket.get("related_ticket_id"):
        preferred_tools.append("lookup_related_ticket")

    if any(
        phrase in text
        for phrase in (
            "pricing",
            "quote",
            "vendor offer",
            "prorating",
            "seat expansion",
            "billing-style",
            "compliance scan",
            "vulnerability",
            "onboarding workflow",
            "permissions error",
            "blocked by an account problem",
        )
    ):
        preferred_tools.append("lookup_internal_routing_note")

    if any(
        phrase in text
        for phrase in ("still", "again", "overdue", "legal", "third update", "priority")
    ):
        preferred_tools.append("lookup_requester_history")

    if infer_ticket_cue(ticket) == "cluster_coordination":
        preferred_tools.append("lookup_queue_cluster_summary")

    if infer_ticket_cue(ticket) == "capacity_planning":
        preferred_tools.append("lookup_queue_capacity_forecast")

    preferred_tools.extend(recommended_tools)

    if hidden_context_remaining:
        preferred_tools.extend(
            [
                "lookup_queue_cluster_summary",
                "lookup_queue_capacity_forecast",
                "lookup_internal_routing_note",
                "lookup_related_ticket",
                "lookup_requester_history",
            ]
        )

    deduped_tools: list[str] = []
    for tool_name in preferred_tools:
        if tool_name not in deduped_tools:
            deduped_tools.append(tool_name)
    return deduped_tools


def select_cue_based_tool(
    ticket: dict[str, Any],
    *,
    hidden_context_remaining: bool,
    used_tools: set[str],
    available_tools: set[str] | None = None,
) -> str | None:
    preferred_tools = preferred_tool_order(
        ticket,
        hidden_context_remaining=hidden_context_remaining,
    )
    available_tool_set = set(available_tools or [])
    for tool_name in preferred_tools:
        if available_tool_set and tool_name not in available_tool_set:
            continue
        if tool_name not in used_tools:
            return tool_name
    return None


def choose_operational_action(
    ticket: dict[str, Any],
    history: list[dict[str, Any]],
    available_action_types: list[str] | None = None,
) -> tuple[HelpdeskTicketAction | None, str | None]:
    if not ticket:
        return None, None
    operational_context = ticket.get("operational_context") or {}
    recommended_actions = list(operational_context.get("recommended_actions") or [])
    available_action_set = set(available_action_types or [])
    current_ticket_id = str(ticket.get("ticket_id", ""))
    prior_ticket_history = [
        entry for entry in history if entry.get("ticket_id") == current_ticket_id
    ]
    used_action_types = {
        entry.get("predicted", {}).get("action_type")
        for entry in prior_ticket_history
        if entry.get("predicted")
    }

    for action_name in ("open_incident", "request_info", "defer"):
        if action_name not in recommended_actions:
            continue
        if available_action_set and action_name not in available_action_set:
            continue
        if action_name in used_action_types:
            continue
        if action_name == "defer" and not ticket.get("tickets_after_current", 0):
            continue
        return HelpdeskTicketAction(action_type=action_name), action_name
    return None, None


def merge_ticket_context(
    ticket: dict[str, Any],
    observation: HelpdeskTicketObservation,
) -> dict[str, Any]:
    merged_ticket = dict(ticket)
    if getattr(observation, "last_tool_result", None) is not None:
        merged_ticket["last_tool_result"] = observation.last_tool_result
        if observation.last_tool_result.get("tool_name") == "lookup_queue_capacity_forecast":
            if observation.last_tool_result.get("future_queue_demand") is not None:
                merged_ticket["future_queue_demand"] = observation.last_tool_result[
                    "future_queue_demand"
                ]
            if observation.last_tool_result.get("capacity_state") is not None:
                merged_ticket["capacity_state"] = observation.last_tool_result[
                    "capacity_state"
                ]
    merged_ticket["recent_history"] = list(getattr(observation, "history", []) or [])
    merged_ticket["queue_position"] = getattr(observation, "queue_position", None)
    merged_ticket["tickets_remaining"] = getattr(observation, "tickets_remaining", None)
    merged_ticket["tickets_after_current"] = getattr(observation, "tickets_after_current", None)
    merged_ticket["available_tools"] = list(getattr(observation, "available_tools", []) or [])
    merged_ticket["available_action_types"] = list(
        getattr(observation, "available_action_types", []) or []
    )
    merged_ticket["last_reward_components"] = dict(
        getattr(observation, "last_reward_components", {}) or {}
    )
    observation_metadata = getattr(observation, "metadata", {}) or {}
    if observation_metadata.get("last_feedback_summary"):
        merged_ticket["feedback_summary"] = observation_metadata["last_feedback_summary"]
    if observation_metadata.get("capacity_state") is not None:
        merged_ticket["capacity_state"] = observation_metadata["capacity_state"]
    if observation_metadata.get("future_queue_demand") is not None:
        merged_ticket["future_queue_demand"] = observation_metadata["future_queue_demand"]
    return merged_ticket


def choose_policy_action(
    policy: PolicyConfig,
    observation: HelpdeskTicketObservation,
    investigations_by_ticket: dict[str, int],
    submit_builder: SubmitBuilder,
    *,
    used_tools_by_ticket: dict[str, set[str]] | None = None,
    adaptive_bandit: AdaptiveToolBandit | None = None,
) -> tuple[HelpdeskTicketAction, str, str | None]:
    ticket = merge_ticket_context(observation.current_ticket or {}, observation)
    ticket_id = str(ticket.get("ticket_id", ""))
    ticket_investigations = investigations_by_ticket.get(ticket_id, 0)
    used_tools = set()
    if used_tools_by_ticket is not None:
        used_tools = set(used_tools_by_ticket.get(ticket_id, set()))
    context_status = ticket.get("context_status") or {}
    hidden_context_remaining = bool(context_status.get("hidden_context_remaining"))
    available_tools = set(getattr(observation, "available_tools", []) or [])

    if ticket_investigations < policy.max_investigations_per_ticket:
        if policy.strategy == "adaptive" and adaptive_bandit is not None and hidden_context_remaining:
            candidate_tools = [
                tool_name
                for tool_name in preferred_tool_order(
                    ticket,
                    hidden_context_remaining=hidden_context_remaining,
                )
                if tool_name not in used_tools and tool_name in available_tools
            ]
            if not candidate_tools:
                candidate_tools = [
                    tool_name
                    for tool_name in AVAILABLE_TOOLS
                    if tool_name not in used_tools and tool_name in available_tools
                ]
            if candidate_tools:
                cue = infer_ticket_cue(ticket)
                tool_name = adaptive_bandit.choose_tool(cue, candidate_tools)
                return (
                    HelpdeskTicketAction(action_type="investigate", tool_name=tool_name),
                    "adaptive_bandit_investigate",
                    cue,
                )

        if policy.investigate_hidden_context and hidden_context_remaining:
            tool_name = select_cue_based_tool(
                ticket,
                hidden_context_remaining=hidden_context_remaining,
                used_tools=used_tools,
                available_tools=available_tools,
            )
            if tool_name is not None:
                return (
                    HelpdeskTicketAction(action_type="investigate", tool_name=tool_name),
                    "investigate_hidden_context",
                    infer_ticket_cue(ticket),
                )
        if (
            policy.investigate_related_ticket_hint
            and ticket.get("related_ticket_id")
            and "lookup_related_ticket" not in used_tools
        ):
            return (
                HelpdeskTicketAction(
                    action_type="investigate",
                    tool_name="lookup_related_ticket",
                ),
                "investigate_related_ticket_hint",
                infer_ticket_cue(ticket),
            )
        if (
            policy.investigate_ambiguity_history
            and (
                ticket.get("ambiguity_note")
                or ticket.get("feedback_summary")
                or hidden_context_remaining
            )
            and "lookup_requester_history" not in used_tools
        ):
            return (
                HelpdeskTicketAction(
                    action_type="investigate",
                    tool_name="lookup_requester_history",
                ),
                "investigate_ambiguity_history",
                infer_ticket_cue(ticket),
            )

    operational_action, operational_source = choose_operational_action(
        ticket,
        list(getattr(observation, "history", []) or []),
        list(getattr(observation, "available_action_types", []) or []),
    )
    if operational_action is not None and operational_source is not None:
        return operational_action, operational_source, infer_ticket_cue(ticket)

    return submit_builder(ticket, list(observation.allowed_fields)), "submit", None


def rollout_episode(
    *,
    env: HelpdeskTicketRoutingEnvironment,
    policy: PolicyConfig,
    seed: int,
    task_id: int,
    submit_builder: SubmitBuilder,
    adaptive_bandit: AdaptiveToolBandit | None = None,
    update_adaptive: bool = False,
) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    task = get_task_definition(task_id)
    observation = env.reset(seed=seed, task_id=task_id)
    investigations_by_ticket: dict[str, int] = {}
    used_tools_by_ticket: dict[str, set[str]] = {}
    episode_return = 0.0
    trajectories: list[dict[str, Any]] = []

    while not observation.done:
        ticket = merge_ticket_context(observation.current_ticket or {}, observation)
        ticket_id = str(ticket.get("ticket_id", ""))
        action, action_source, action_cue = choose_policy_action(
            policy,
            observation,
            investigations_by_ticket,
            submit_builder,
            used_tools_by_ticket=used_tools_by_ticket,
            adaptive_bandit=adaptive_bandit,
        )
        next_observation = env.step(action)
        reward_value = float(next_observation.reward or 0.0)
        episode_return += reward_value
        if action.action_type == "investigate" and ticket_id:
            investigations_by_ticket[ticket_id] = investigations_by_ticket.get(ticket_id, 0) + 1
            used_tools_by_ticket.setdefault(ticket_id, set()).add(str(action.tool_name))
            if policy.strategy == "adaptive" and adaptive_bandit is not None and update_adaptive:
                adaptive_bandit.record_reward(
                    action_cue or infer_ticket_cue(ticket),
                    str(action.tool_name),
                    reward_value,
                )

        history_entry = env.state.history_entries[-1] if env.state.history_entries else {}
        trajectories.append(
            {
                "policy": policy.name,
                "seed": seed,
                "task_id": task_id,
                "task_name": task["name"],
                "episode_id": env.state.episode_id,
                "step_index": len(trajectories) + 1,
                "ticket_id": history_entry.get("ticket_id", ticket_id),
                "action_source": action_source,
                "action_cue": action_cue,
                "action": action.model_dump(exclude_none=True),
                "step_reward": reward_value,
                "rubric_reward": next_observation.rubric_reward,
                "done": next_observation.done,
                "feedback_summary": history_entry.get("feedback_summary"),
                "reward_kind": history_entry.get("reward_kind"),
                "score": history_entry.get("score"),
                "breakdown": history_entry.get("breakdown", {}),
                "reward_components": history_entry.get("reward_components", {}),
                "context_status_before_action": ticket.get("context_status"),
            }
        )
        observation = next_observation

    queue_size = max(1, len(env.state.queue_ticket_ids))
    terminal_reward = float(observation.reward or 0.0)
    terminal_rubric_reward = (
        float(observation.rubric_reward)
        if observation.rubric_reward is not None
        else terminal_reward
    )
    summary = {
        "policy": policy.name,
        "policy_config": asdict(policy),
        "seed": seed,
        "task_id": task_id,
        "task_name": task["name"],
        "episode_id": env.state.episode_id,
        "queue_size": queue_size,
        "step_count": env.state.step_count,
        "tickets_processed": len(env.state.per_ticket_scores),
        "investigation_steps": env.state.investigation_steps,
        "episode_return": episode_return,
        "normalized_return": episode_return / queue_size,
        "terminal_reward": terminal_reward,
        "terminal_rubric_reward": terminal_rubric_reward,
        "average_ticket_score": env.state.average_score_so_far,
        "queue_management_score": env.state.queue_management_score,
        "planning_penalty_total": env.state.planning_penalty_total,
        "capacity_pressure_tickets_resolved": env.state.capacity_pressure_tickets_resolved,
        "per_ticket_scores": list(env.state.per_ticket_scores),
    }
    if adaptive_bandit is not None and policy.strategy == "adaptive":
        summary["learned_tool_values"] = adaptive_bandit.export()
    return summary, trajectories


def _safe_mean(values: list[float]) -> float:
    if not values:
        return 0.0
    return round(mean(values), 6)


def summarize_policy_episodes(
    policy: PolicyConfig,
    episode_summaries: list[dict[str, Any]],
) -> dict[str, Any]:
    per_task: dict[str, Any] = {}
    for task_id in TASK_IDS:
        task_episodes = [
            episode for episode in episode_summaries if episode["task_id"] == task_id
        ]
        if not task_episodes:
            continue
        per_task[str(task_id)] = {
            "episodes": len(task_episodes),
            "avg_episode_return": _safe_mean(
                [float(episode["episode_return"]) for episode in task_episodes]
            ),
            "avg_normalized_return": _safe_mean(
                [float(episode["normalized_return"]) for episode in task_episodes]
            ),
            "avg_terminal_reward": _safe_mean(
                [float(episode["terminal_reward"]) for episode in task_episodes]
            ),
            "avg_terminal_rubric_reward": _safe_mean(
                [float(episode["terminal_rubric_reward"]) for episode in task_episodes]
            ),
            "avg_queue_management_score": _safe_mean(
                [float(episode["queue_management_score"]) for episode in task_episodes]
            ),
            "avg_planning_penalty_total": _safe_mean(
                [float(episode["planning_penalty_total"]) for episode in task_episodes]
            ),
            "avg_capacity_pressure_tickets_resolved": _safe_mean(
                [
                    float(episode["capacity_pressure_tickets_resolved"])
                    for episode in task_episodes
                ]
            ),
            "avg_investigation_steps": _safe_mean(
                [float(episode["investigation_steps"]) for episode in task_episodes]
            ),
        }

    return {
        "policy": policy.name,
        "config": asdict(policy),
        "episodes": len(episode_summaries),
        "avg_episode_return": _safe_mean(
            [float(episode["episode_return"]) for episode in episode_summaries]
        ),
        "avg_normalized_return": _safe_mean(
            [float(episode["normalized_return"]) for episode in episode_summaries]
        ),
        "avg_terminal_reward": _safe_mean(
            [float(episode["terminal_reward"]) for episode in episode_summaries]
        ),
        "avg_terminal_rubric_reward": _safe_mean(
            [float(episode["terminal_rubric_reward"]) for episode in episode_summaries]
        ),
        "avg_queue_management_score": _safe_mean(
            [float(episode["queue_management_score"]) for episode in episode_summaries]
        ),
        "avg_planning_penalty_total": _safe_mean(
            [float(episode["planning_penalty_total"]) for episode in episode_summaries]
        ),
        "avg_capacity_pressure_tickets_resolved": _safe_mean(
            [
                float(episode["capacity_pressure_tickets_resolved"])
                for episode in episode_summaries
            ]
        ),
        "avg_investigation_steps": _safe_mean(
            [float(episode["investigation_steps"]) for episode in episode_summaries]
        ),
        "avg_ticket_score": _safe_mean(
            [float(episode["average_ticket_score"]) for episode in episode_summaries]
        ),
        "per_task": per_task,
    }


def evaluate_policy(
    policy: PolicyConfig,
    seeds: Iterable[int],
    task_ids: Iterable[int],
    *,
    env_factory: EnvFactory = HelpdeskTicketRoutingEnvironment,
    submit_builder: SubmitBuilder = default_submit_builder,
    adaptive_bandit: AdaptiveToolBandit | None = None,
    update_adaptive: bool = False,
) -> dict[str, Any]:
    episode_summaries: list[dict[str, Any]] = []
    trajectories: list[dict[str, Any]] = []

    for seed in seeds:
        for task_id in task_ids:
            env = env_factory()
            summary, episode_trajectories = rollout_episode(
                env=env,
                policy=policy,
                seed=seed,
                task_id=task_id,
                submit_builder=submit_builder,
                adaptive_bandit=adaptive_bandit,
                update_adaptive=update_adaptive,
            )
            episode_summaries.append(summary)
            trajectories.extend(episode_trajectories)

    result = {
        "policy": policy.name,
        "summary": summarize_policy_episodes(policy, episode_summaries),
        "episodes": episode_summaries,
        "trajectories": trajectories,
    }
    if adaptive_bandit is not None and policy.strategy == "adaptive":
        result["adaptive_bandit"] = adaptive_bandit.export()
    return result


def _selection_tuple(summary: dict[str, Any]) -> tuple[float, float, float, float, float, float]:
    return (
        float(summary["avg_terminal_rubric_reward"]),
        float(summary["avg_queue_management_score"]),
        -float(summary["avg_planning_penalty_total"]),
        float(summary["avg_episode_return"]),
        float(summary["avg_normalized_return"]),
        -float(summary["avg_investigation_steps"]),
    )


def select_best_policy(policy_runs: list[dict[str, Any]]) -> dict[str, Any]:
    return max(policy_runs, key=lambda run: _selection_tuple(run["summary"]))


def _delta(best: dict[str, Any], baseline: dict[str, Any], key: str) -> float:
    return round(float(best[key]) - float(baseline[key]), 6)


def _write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def _write_jsonl(path: Path, records: Iterable[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, sort_keys=True) + "\n")


def compare_policies(
    policies: list[PolicyConfig],
    seeds: list[int],
    task_ids: list[int],
    *,
    output_dir: Path,
    env_factory: EnvFactory = HelpdeskTicketRoutingEnvironment,
    submit_builder: SubmitBuilder = default_submit_builder,
) -> dict[str, Any]:
    output_dir = Path(output_dir)
    policy_runs = []
    for policy in policies:
        adaptive_bandit = AdaptiveToolBandit() if policy.strategy == "adaptive" else None
        policy_runs.append(
            evaluate_policy(
                policy,
                seeds,
                task_ids,
                env_factory=env_factory,
                submit_builder=submit_builder,
                adaptive_bandit=adaptive_bandit,
                update_adaptive=policy.strategy == "adaptive",
            )
        )
    best_run = select_best_policy(policy_runs)
    baseline_run = policy_runs[0]

    report = {
        "mode": "compare",
        "task_ids": task_ids,
        "seeds": seeds,
        "selection_metric": (
            "avg_terminal_rubric_reward_then_queue_management_then_lower_planning_penalty"
        ),
        "baseline_policy": baseline_run["policy"],
        "best_policy": best_run["policy"],
        "improvement_vs_baseline": {
            "avg_episode_return": _delta(
                best_run["summary"], baseline_run["summary"], "avg_episode_return"
            ),
            "avg_normalized_return": _delta(
                best_run["summary"], baseline_run["summary"], "avg_normalized_return"
            ),
            "avg_terminal_reward": _delta(
                best_run["summary"], baseline_run["summary"], "avg_terminal_reward"
            ),
            "avg_terminal_rubric_reward": _delta(
                best_run["summary"],
                baseline_run["summary"],
                "avg_terminal_rubric_reward",
            ),
            "avg_queue_management_score": _delta(
                best_run["summary"],
                baseline_run["summary"],
                "avg_queue_management_score",
            ),
            "avg_planning_penalty_total": _delta(
                best_run["summary"],
                baseline_run["summary"],
                "avg_planning_penalty_total",
            ),
        },
        "policy_summaries": [run["summary"] for run in policy_runs],
        "ranking": [
            run["policy"]
            for run in sorted(
                policy_runs,
                key=lambda run: _selection_tuple(run["summary"]),
                reverse=True,
            )
        ],
        "adaptive_bandits": {
            run["policy"]: run["adaptive_bandit"]
            for run in policy_runs
            if "adaptive_bandit" in run
        },
        "artifacts": {
            "summary": str(output_dir / "compare_summary.json"),
            "episodes": str(output_dir / "compare_episodes.jsonl"),
            "trajectories": str(output_dir / "compare_trajectories.jsonl"),
        },
    }

    _write_json(output_dir / "compare_summary.json", report)
    _write_jsonl(
        output_dir / "compare_episodes.jsonl",
        (
            {"policy": run["policy"], **episode}
            for run in policy_runs
            for episode in run["episodes"]
        ),
    )
    _write_jsonl(
        output_dir / "compare_trajectories.jsonl",
        (trajectory for run in policy_runs for trajectory in run["trajectories"]),
    )
    return report


def search_policies(
    candidate_policies: list[PolicyConfig],
    train_seeds: list[int],
    eval_seeds: list[int],
    task_ids: list[int],
    *,
    output_dir: Path,
    env_factory: EnvFactory = HelpdeskTicketRoutingEnvironment,
    submit_builder: SubmitBuilder = default_submit_builder,
    baseline_policy_name: str = "no_investigation",
) -> dict[str, Any]:
    output_dir = Path(output_dir)
    train_runs = []
    trained_bandits: dict[str, AdaptiveToolBandit] = {}
    for policy in candidate_policies:
        adaptive_bandit = AdaptiveToolBandit() if policy.strategy == "adaptive" else None
        train_run = evaluate_policy(
            policy,
            train_seeds,
            task_ids,
            env_factory=env_factory,
            submit_builder=submit_builder,
            adaptive_bandit=adaptive_bandit,
            update_adaptive=policy.strategy == "adaptive",
        )
        train_runs.append(train_run)
        if adaptive_bandit is not None:
            trained_bandits[policy.name] = adaptive_bandit.frozen_copy()
    selected_run = select_best_policy(train_runs)
    selected_policy = POLICY_LIBRARY[selected_run["policy"]]
    eval_selected = evaluate_policy(
        selected_policy,
        eval_seeds,
        task_ids,
        env_factory=env_factory,
        submit_builder=submit_builder,
        adaptive_bandit=trained_bandits.get(selected_policy.name),
        update_adaptive=False,
    )

    baseline_policy = POLICY_LIBRARY.get(baseline_policy_name, candidate_policies[0])
    eval_baseline = evaluate_policy(
        baseline_policy,
        eval_seeds,
        task_ids,
        env_factory=env_factory,
        submit_builder=submit_builder,
        adaptive_bandit=trained_bandits.get(baseline_policy.name),
        update_adaptive=False,
    )

    report = {
        "mode": "search",
        "task_ids": task_ids,
        "train_seeds": train_seeds,
        "eval_seeds": eval_seeds,
        "selection_metric": (
            "avg_terminal_rubric_reward_then_queue_management_then_lower_planning_penalty"
        ),
        "candidate_policies": [policy.name for policy in candidate_policies],
        "selected_policy": selected_policy.name,
        "baseline_policy": baseline_policy.name,
        "train_policy_summaries": [run["summary"] for run in train_runs],
        "trained_adaptive_bandits": {
            name: bandit.export() for name, bandit in trained_bandits.items()
        },
        "eval_selected_summary": eval_selected["summary"],
        "eval_baseline_summary": eval_baseline["summary"],
        "eval_improvement_vs_baseline": {
            "avg_episode_return": _delta(
                eval_selected["summary"],
                eval_baseline["summary"],
                "avg_episode_return",
            ),
            "avg_normalized_return": _delta(
                eval_selected["summary"],
                eval_baseline["summary"],
                "avg_normalized_return",
            ),
            "avg_terminal_reward": _delta(
                eval_selected["summary"],
                eval_baseline["summary"],
                "avg_terminal_reward",
            ),
            "avg_terminal_rubric_reward": _delta(
                eval_selected["summary"],
                eval_baseline["summary"],
                "avg_terminal_rubric_reward",
            ),
            "avg_queue_management_score": _delta(
                eval_selected["summary"],
                eval_baseline["summary"],
                "avg_queue_management_score",
            ),
            "avg_planning_penalty_total": _delta(
                eval_selected["summary"],
                eval_baseline["summary"],
                "avg_planning_penalty_total",
            ),
        },
        "artifacts": {
            "summary": str(output_dir / "search_summary.json"),
            "train_episodes": str(output_dir / "search_train_episodes.jsonl"),
            "train_trajectories": str(output_dir / "search_train_trajectories.jsonl"),
            "eval_episodes": str(output_dir / "search_eval_episodes.jsonl"),
            "eval_trajectories": str(output_dir / "search_eval_trajectories.jsonl"),
        },
    }

    _write_json(output_dir / "search_summary.json", report)
    _write_jsonl(
        output_dir / "search_train_episodes.jsonl",
        (
            {"policy": run["policy"], **episode}
            for run in train_runs
            for episode in run["episodes"]
        ),
    )
    _write_jsonl(
        output_dir / "search_train_trajectories.jsonl",
        (trajectory for run in train_runs for trajectory in run["trajectories"]),
    )
    _write_jsonl(
        output_dir / "search_eval_episodes.jsonl",
        (
            {"policy": eval_selected["policy"], **episode}
            for episode in eval_selected["episodes"]
        ),
    )
    _write_jsonl(
        output_dir / "search_eval_trajectories.jsonl",
        (trajectory for trajectory in eval_selected["trajectories"]),
    )
    return report


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Run seeded local rollouts and a small policy-improvement loop for the "
            "IT helpdesk OpenEnv environment."
        )
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    compare_parser = subparsers.add_parser(
        "compare",
        help="Compare fixed policy choices across repeated seeded rollouts.",
    )
    compare_parser.add_argument(
        "--policies",
        default=",".join(DEFAULT_COMPARE_POLICIES),
        help=f"Comma-separated policy names. Available: {', '.join(POLICY_LIBRARY)}",
    )
    compare_parser.add_argument(
        "--seeds",
        default="42-51",
        help="Comma-separated seeds or ranges, for example 42-51 or 42,50,60.",
    )
    compare_parser.add_argument(
        "--task-ids",
        default="1,2,3",
        help="Comma-separated task IDs or ranges, for example 1,2,3 or 1-3.",
    )
    compare_parser.add_argument(
        "--output-dir",
        default=DEFAULT_OUTPUT_DIR,
        help="Directory for JSON and JSONL artifacts.",
    )

    search_parser = subparsers.add_parser(
        "search",
        help="Select the best policy on train seeds, then re-evaluate on holdout seeds.",
    )
    search_parser.add_argument(
        "--candidate-policies",
        default=",".join(DEFAULT_SEARCH_POLICIES),
        help=f"Comma-separated candidate policy names. Available: {', '.join(POLICY_LIBRARY)}",
    )
    search_parser.add_argument(
        "--train-seeds",
        default="40-49",
        help="Train seeds used for reward-based policy selection.",
    )
    search_parser.add_argument(
        "--eval-seeds",
        default="50-59",
        help="Holdout seeds used for the selected policy evaluation.",
    )
    search_parser.add_argument(
        "--task-ids",
        default="1,2,3",
        help="Comma-separated task IDs or ranges, for example 1,2,3 or 1-3.",
    )
    search_parser.add_argument(
        "--baseline-policy",
        default="no_investigation",
        help="Baseline policy used for the final improvement delta.",
    )
    search_parser.add_argument(
        "--output-dir",
        default=DEFAULT_OUTPUT_DIR,
        help="Directory for JSON and JSONL artifacts.",
    )

    return parser


def _print_summary(label: str, summary: dict[str, Any]) -> None:
    print(
        json.dumps(
            {
                label: {
                    "policy": summary["policy"],
                    "avg_episode_return": summary["avg_episode_return"],
                    "avg_normalized_return": summary["avg_normalized_return"],
                    "avg_terminal_reward": summary["avg_terminal_reward"],
                    "avg_terminal_rubric_reward": summary["avg_terminal_rubric_reward"],
                    "avg_planning_penalty_total": summary["avg_planning_penalty_total"],
                    "avg_investigation_steps": summary["avg_investigation_steps"],
                }
            },
            sort_keys=True,
        )
    )


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()

    output_dir = Path(args.output_dir)

    if args.command == "compare":
        policies = resolve_policies(args.policies)
        seeds = parse_int_spec(args.seeds, field_name="seeds")
        task_ids = parse_task_ids(args.task_ids)
        report = compare_policies(
            policies,
            seeds,
            task_ids,
            output_dir=output_dir,
        )
        print(json.dumps(report, indent=2, sort_keys=True))
        return

    candidate_policies = resolve_policies(args.candidate_policies)
    train_seeds = parse_int_spec(args.train_seeds, field_name="train_seeds")
    eval_seeds = parse_int_spec(args.eval_seeds, field_name="eval_seeds")
    task_ids = parse_task_ids(args.task_ids)
    report = search_policies(
        candidate_policies,
        train_seeds,
        eval_seeds,
        task_ids,
        output_dir=output_dir,
        baseline_policy_name=args.baseline_policy,
    )
    print(json.dumps(report, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()