Add policy learning loop and strengthen RL-style environment

.gitignore

@@ -6,3 +6,7 @@ __pycache__/
 .mypy_cache/
 .ruff_cache/
 build/
+analysis/policy_learning_runs/
+analysis/policy_learning_test/
+analysis/policy_learning_compare_test/
+analysis/policy_learning_runs_smoke/

README.md

@@ -38,6 +38,8 @@ The environment models a realistic helpdesk workflow:
 4. the grader assigns deterministic credit
 5. the environment advances to the next ticket until the queue is complete
 
+For hard-task tickets, the environment can now withhold decisive routing context until the agent uses the right investigation tool. That keeps the task from collapsing into one-shot classification and makes tool choice part of the policy.
+
 This domain is useful for OpenEnv because it is operationally realistic, easy to evaluate with typed outputs, and naturally supports a clean easy-to-hard task ladder.
 
 ## Why This Is A Good Hackathon Domain
@@ -59,6 +61,37 @@ The project uses a queue-based episode model.
 
 The environment classes and vocabulary are intentionally frozen to keep collaboration and judging simple.
 
+## Lightweight Policy Improvement Loop
+
+The repo now includes a small local learning runner in `policy_learning.py`. It does not update model weights, but it does run repeated rollouts over many seeds, log full trajectories, and select the best policy configuration from a discrete candidate set using observed reward.
+
+That gives the project a real improvement loop for judge demos:
+
+- compare `no_investigation` against `investigate_when_context_hidden`
+- log per-step rewards, feedback summaries, and reward components to JSONL
+- search over small policy variants such as `legacy_single_probe`, `context_chain`, and `hybrid_context`
+- select the best policy on train seeds, then re-evaluate it on holdout seeds
+
+Example commands:
+
+```bash
+python policy_learning.py compare --seeds 42-51 --task-ids 1,2,3
+python policy_learning.py search --train-seeds 40-49 --eval-seeds 50-59 --task-ids 1,2,3
+```
+
+Artifacts are written to `analysis/policy_learning_runs/` by default:
+
+- `compare_summary.json`
+- `compare_episodes.jsonl`
+- `compare_trajectories.jsonl`
+- `search_summary.json`
+- `search_train_episodes.jsonl`
+- `search_train_trajectories.jsonl`
+- `search_eval_episodes.jsonl`
+- `search_eval_trajectories.jsonl`
+
+The default submit policy inside this runner stays deterministic and local. It reuses the repo's heuristic routing logic, so the discrete policy search focuses on investigation behavior and reward-driven policy selection rather than on external LLM latency or API cost.
+
 ## Task Ladder
 
 | ID | Name | Difficulty | Required Fields | What The Agent Must Do |
@@ -125,6 +158,7 @@ Each observation also includes:
 - `task_name`
 - `instructions`
 - `allowed_fields`
+- `available_action_types`
 - `available_tools`
 - `investigation_budget_remaining`
 - `last_tool_result`
@@ -133,7 +167,12 @@ Each observation also includes:
 - `tickets_after_current`
 - `tickets_processed`
 - `queue_position`
+- `average_score_so_far`
+- `progress_fraction`
 - `history`
+- `last_reward_components`
+- `rubric_reward` on terminal observations
+- `metadata.last_feedback_summary` for compact reward / penalty feedback
 - standard OpenEnv fields such as `done` and `reward`
 
 The internal `HelpdeskTicketState` tracks:
@@ -162,6 +201,15 @@ Available tools:
 
 - `lookup_related_ticket`
 - `lookup_requester_history`
+- `lookup_internal_routing_note`
+
+Hard-task investigation behavior:
+
+- some ambiguous and non-default-routing tickets start with redacted descriptions
+- linked-ticket previews and internal routing notes stay hidden until the matching tool is used
+- useful investigation steps return a small positive shaping reward
+- premature hard-task submission can incur a shaping penalty even when the visible text looks plausible
+- terminal `rubric_reward` remains the objective evaluation signal, while per-step `reward` is the denser training signal
 
 Per-field behavior:
 
@@ -190,6 +238,12 @@ Step reward is lightly milestone-shaped: high per-ticket scores get a small bonu
 
 Final reward also includes a tiny queue-economics penalty only when the agent exceeds the free investigation budget. One investigation per queued ticket is free; extra investigation steps reduce the final reward slightly.
 
+To make the environment more RL-friendly, each observation now also surfaces structured reward telemetry:
+
+- `last_reward_components` exposes ticket score, shaped step reward, milestone adjustment, trajectory reward when applicable, and any investigation penalty applied
+- `average_score_so_far` and `progress_fraction` expose trajectory progress without leaking future labels
+- `history` retains the same reward components plus a compact `feedback_summary` string for downstream agents
+
 ## Grounded Scoring
 
 The grader is intentionally not fuzzy by default.
@@ -343,6 +397,7 @@ Optional target:
 
 - `ENV_URL`
 - default value: `http://localhost:7860`
+- `SEED`
 - `TASK_ID`
 - `RUN_ALL_TASKS`
 

inference.py

@@ -66,13 +66,27 @@ from vocabulary import (
 DEFAULT_API_BASE_URL = "https://router.huggingface.co/v1"
 DEFAULT_MODEL_NAME = "<your-active-model>"
 
+
+def _get_int_env(name: str, default: int) -> int:
+    raw_value = os.getenv(name)
+    if raw_value is None or raw_value.strip() == "":
+        return default
+    try:
+        return int(raw_value)
+    except ValueError:
+        print(
+            f"[WARN] {name}={raw_value!r} is not a valid integer; using {default}.",
+            flush=True,
+        )
+        return default
+
 API_BASE_URL = os.getenv("API_BASE_URL", DEFAULT_API_BASE_URL)
 MODEL_NAME = os.getenv("MODEL_NAME", DEFAULT_MODEL_NAME)
 HF_TOKEN = os.getenv("HF_TOKEN")
 LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
 ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")
 
-SEED = 42
+SEED = _get_int_env("SEED", 42)
 TASK_ID_ENV = os.getenv("TASK_ID")
 RUN_ALL_TASKS_ENV = os.getenv("RUN_ALL_TASKS", "").strip().lower() in {
     "1",
@@ -94,6 +108,14 @@ if llm_mode_enabled():
     llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
 
 
+RECENT_HISTORY_LIMIT = 2
+ROUTING_PRIORS = "\n".join(
+    f"- {issue_type}: assignment_group={ISSUE_TYPE_TO_ASSIGNMENT_GROUP[issue_type]}, "
+    f"resolution_action={ISSUE_TYPE_TO_RESOLUTION_ACTION[issue_type]}"
+    for issue_type in ISSUE_TYPES
+)
+
+
 SYSTEM_PROMPT = """\
 You are an expert IT helpdesk ticket routing agent. Given a helpdesk ticket, you must produce a JSON object with the requested fields.
 
@@ -103,19 +125,79 @@ Valid values:
 - assignment_group: {assignment_groups}
 - resolution_action: {resolution_actions}
 
+Decision rules:
+- Follow this environment's label ontology exactly; do not invent categories.
+- Prefer the primary operational workflow label over a secondary technical symptom.
+- Keep assignment_group and resolution_action consistent with the chosen issue_type unless the ticket explicitly justifies a different choice.
+- Use investigation results and recent evaluation feedback when provided.
+
+Domain conventions:
+- Enterprise pricing, quotes, plan comparisons, and commercial procurement requests map to service_request, usually with medium priority.
+- Onboarding work that is blocked by an access problem still maps to onboarding when the primary workflow is onboarding; the assignment_group may still be service_desk if the ticket says onboarding cannot resolve the access issue.
+- Single-user sign-in, login, MFA, or 2FA lockouts map to identity_access and are usually high priority, not critical.
+- Reserve critical priority for outages, widespread business blockers, or explicit urgent critical incidents.
+
+Routing priors:
+{routing_priors}
+
 Return ONLY valid JSON with the requested fields. No markdown, no explanation.""".format(
     issue_types=", ".join(ISSUE_TYPES),
     priorities=", ".join(PRIORITIES),
     assignment_groups=", ".join(ASSIGNMENT_GROUPS),
     resolution_actions=", ".join(RESOLUTION_ACTIONS),
+    routing_priors=ROUTING_PRIORS,
 )
 
 
-def call_llm(ticket: dict, allowed_fields: list[str], instructions: str) -> dict:
-    assert llm_client is not None, "LLM client not configured"
+def format_recent_history_entries(
+    history: list[dict[str, Any]], limit: int = RECENT_HISTORY_LIMIT
+) -> str:
+    if not history:
+        return ""
+
+    lines = ["Recent evaluation feedback (latest last):"]
+    for entry in history[-limit:]:
+        predicted = json.dumps(entry.get("predicted", {}), sort_keys=True)
+        line = (
+            f"- Ticket {entry.get('ticket_id', '?')}: predicted={predicted}, "
+            f"score={entry.get('score', 0.0)}"
+        )
+        feedback_summary = entry.get("feedback_summary")
+        if feedback_summary:
+            line += f", feedback={feedback_summary}"
+        reward = entry.get("reward")
+        if reward is not None:
+            line += f", reward={reward}"
+        rubric_reward = entry.get("rubric_reward")
+        if rubric_reward is not None:
+            line += f", rubric_reward={rubric_reward}"
+        breakdown = entry.get("breakdown") or {}
+        if breakdown:
+            line += f", breakdown={json.dumps(breakdown, sort_keys=True)}"
+        penalty_reason = entry.get("penalty_reason")
+        if penalty_reason:
+            line += f", penalty_reason={penalty_reason}"
+        tool_result = entry.get("tool_result")
+        if tool_result is not None:
+            line += f", tool_result={json.dumps(tool_result, sort_keys=True)}"
+        reward_components = entry.get("reward_components")
+        if reward_components:
+            line += f", reward_components={json.dumps(reward_components, sort_keys=True)}"
+        lines.append(line)
+    return "\n".join(lines)
+
+
+def build_llm_user_message(ticket: dict, allowed_fields: list[str], instructions: str) -> str:
     ambiguity_note = ticket.get("ambiguity_note")
     related_preview = ticket.get("related_ticket_preview") or {}
     last_tool_result = ticket.get("last_tool_result")
+    context_status = ticket.get("context_status") or {}
+    recent_history = ticket.get("recent_history") or []
+    feedback_summary = ticket.get("feedback_summary")
+    last_reward_components = ticket.get("last_reward_components") or {}
+    investigation_budget_remaining = ticket.get("investigation_budget_remaining")
+    average_score_so_far = ticket.get("average_score_so_far")
+    progress_fraction = ticket.get("progress_fraction")
     extra_context_lines: list[str] = []
     if ambiguity_note:
         extra_context_lines.append(f"Ambiguity note: {ambiguity_note}")
@@ -132,20 +214,53 @@ def call_llm(ticket: dict, allowed_fields: list[str], instructions: str) -> dict
         extra_context_lines.append(
             "Investigation result: " + json.dumps(last_tool_result, sort_keys=True)
         )
+    if context_status:
+        extra_context_lines.append(
+            "Context status: " + json.dumps(context_status, sort_keys=True)
+        )
+    if feedback_summary:
+        extra_context_lines.append(f"Latest environment feedback: {feedback_summary}")
+    if last_reward_components:
+        extra_context_lines.append(
+            "Latest reward components: "
+            + json.dumps(last_reward_components, sort_keys=True)
+        )
+    recent_history_block = format_recent_history_entries(recent_history)
+    if recent_history_block:
+        extra_context_lines.append(recent_history_block)
+    queue_position = ticket.get("queue_position")
+    tickets_remaining = ticket.get("tickets_remaining")
+    if queue_position is not None and tickets_remaining is not None:
+        extra_context_lines.append(
+            f"Queue context: queue_position={queue_position}, tickets_remaining={tickets_remaining}"
+        )
+    if average_score_so_far is not None:
+        extra_context_lines.append(f"Average score so far: {average_score_so_far}")
+    if progress_fraction is not None:
+        extra_context_lines.append(f"Episode progress: {progress_fraction}")
+    if investigation_budget_remaining is not None:
+        extra_context_lines.append(
+            f"Investigation budget remaining: {investigation_budget_remaining}"
+        )
     extra_context_block = ""
     if extra_context_lines:
         extra_context_block = "\n" + "\n".join(extra_context_lines)
 
-    user_msg = (
+    return (
         f"Instructions: {instructions}\n\n"
         f"Allowed fields: {', '.join(allowed_fields)}\n\n"
-        f"Title: {ticket['title']}\n"
-        f"Requester: {ticket['requester']}\n"
-        f"Description: {ticket['description']}"
+        f"Title: {ticket.get('title', '')}\n"
+        f"Requester: {ticket.get('requester', '')}\n"
+        f"Description: {ticket.get('description', '')}"
         f"{extra_context_block}\n\n"
         f"Respond with JSON containing ONLY these fields: {', '.join(allowed_fields)}"
     )
 
+
+def call_llm(ticket: dict, allowed_fields: list[str], instructions: str) -> dict:
+    assert llm_client is not None, "LLM client not configured"
+    user_msg = build_llm_user_message(ticket, allowed_fields, instructions)
+
     response = llm_client.chat.completions.create(
         model=MODEL_NAME,
         messages=[
@@ -298,6 +413,95 @@ FULFILL_KEYWORDS = (
     "mfa enabled",
 )
 
+PRICING_REQUEST_KEYWORDS = (
+    "pricing breakdown",
+    "enterprise tier pricing",
+    "enterprise plan",
+    "compare your enterprise plan",
+    "comparing your enterprise plan",
+    "quote",
+    "pricing quote",
+    "commercial proposal",
+    "vendor comparison",
+)
+
+ONBOARDING_WORKFLOW_KEYWORDS = (
+    "onboarding",
+    "new hire",
+    "contractor",
+    "provisioned",
+    "kickoff onboarding",
+)
+
+ACCESS_BLOCKER_KEYWORDS = (
+    "access issue",
+    "permissions error",
+    "permission error",
+    "account access is blocked",
+    "cannot sign in",
+    "can't sign in",
+    "locked",
+    "2fa",
+    "mfa",
+)
+
+SERVICE_DESK_ONBOARDING_ESCALATION_KEYWORDS = (
+    "onboarding team cannot resolve access issues",
+    "routing to service desk",
+    "route to service desk",
+    "service desk",
+)
+
+CRITICAL_INCIDENT_KEYWORDS = (
+    "outage",
+    "company-wide",
+    "all users",
+    "widespread",
+    "production down",
+    "critical incident",
+    "sev1",
+)
+
+HIGH_PRIORITY_SIGNAL_KEYWORDS = (
+    "locked",
+    "blocked",
+    "cannot sign in",
+    "can't sign in",
+    "2fa",
+    "mfa",
+    "expedite",
+    "start monday",
+    "asap",
+    "today",
+    "eod",
+    "urgent",
+)
+
+TIME_SENSITIVE_PRIORITY_KEYWORDS = (
+    "expedite",
+    "start monday",
+    "today",
+    "asap",
+    "eod",
+    "urgent",
+    "immediately",
+)
+
+
+def build_routing_text(ticket: dict) -> str:
+    related_preview = ticket.get("related_ticket_preview") or {}
+    last_tool_result = ticket.get("last_tool_result") or {}
+    return " ".join(
+        [
+            ticket.get("title", ""),
+            ticket.get("description", ""),
+            ticket.get("ambiguity_note", ""),
+            related_preview.get("title", ""),
+            related_preview.get("description", ""),
+            json.dumps(last_tool_result, sort_keys=True),
+        ]
+    ).lower()
+
 
 def heuristic_priority(text: str) -> str:
     if any(word in text for word in CRITICAL_PRIORITY_KEYWORDS):
@@ -323,26 +527,32 @@ def heuristic_resolution_action(text: str, issue_type: str) -> str:
     return ISSUE_TYPE_TO_RESOLUTION_ACTION.get(issue_type, "acknowledge")
 
 
-def heuristic_action(ticket: dict, allowed_fields: list[str]) -> dict:
-    related_preview = ticket.get("related_ticket_preview") or {}
-    last_tool_result = ticket.get("last_tool_result") or {}
-    text = " ".join(
-        [
-            ticket.get("title", ""),
-            ticket.get("description", ""),
-            ticket.get("ambiguity_note", ""),
-            related_preview.get("title", ""),
-            related_preview.get("description", ""),
-            json.dumps(last_tool_result, sort_keys=True),
-        ]
-    ).lower()
+def heuristic_assignment_group(text: str, issue_type: str) -> str:
+    if issue_type == "onboarding":
+        if any(keyword in text for keyword in SERVICE_DESK_ONBOARDING_ESCALATION_KEYWORDS):
+            return "service_desk"
+        if any(keyword in text for keyword in ACCESS_BLOCKER_KEYWORDS) and any(
+            keyword in text for keyword in ONBOARDING_WORKFLOW_KEYWORDS
+        ):
+            return "service_desk"
+    return ISSUE_TYPE_TO_ASSIGNMENT_GROUP.get(issue_type, "service_desk")
 
+
+def infer_issue_type(text: str) -> str:
     issue_type = "general_inquiry"
     for kw, mapped_issue_type in KEYWORD_ISSUE_TYPES.items():
         if kw in text:
             issue_type = mapped_issue_type
             break
+    return issue_type
+
+
+def heuristic_action(
+    ticket: dict, allowed_fields: list[str], issue_type_override: str | None = None
+) -> dict:
+    text = build_routing_text(ticket)
 
+    issue_type = issue_type_override or infer_issue_type(text)
     priority = heuristic_priority(text)
     resolution_action = heuristic_resolution_action(text, issue_type)
 
@@ -352,14 +562,75 @@ def heuristic_action(ticket: dict, allowed_fields: list[str]) -> dict:
     if "priority" in allowed_fields:
         result["priority"] = priority
     if "assignment_group" in allowed_fields:
-        result["assignment_group"] = ISSUE_TYPE_TO_ASSIGNMENT_GROUP.get(
-            issue_type, "service_desk"
-        )
+        result["assignment_group"] = heuristic_assignment_group(text, issue_type)
     if "resolution_action" in allowed_fields:
         result["resolution_action"] = resolution_action
     return result
 
 
+def apply_domain_overrides(
+    ticket: dict, candidate: dict[str, Any], allowed_fields: list[str]
+) -> tuple[dict[str, Any], list[str]]:
+    updated = dict(candidate)
+    reasons: list[str] = []
+    text = build_routing_text(ticket)
+
+    issue_type = updated.get("issue_type")
+    if "issue_type" in allowed_fields and issue_type is not None:
+        if (
+            issue_type in {"billing_license", "general_inquiry"}
+            and any(keyword in text for keyword in PRICING_REQUEST_KEYWORDS)
+        ):
+            updated["issue_type"] = "service_request"
+            issue_type = "service_request"
+            reasons.append("override_issue_type=service_request(pricing_request)")
+        elif (
+            issue_type == "identity_access"
+            and any(keyword in text for keyword in ONBOARDING_WORKFLOW_KEYWORDS)
+            and any(keyword in text for keyword in ACCESS_BLOCKER_KEYWORDS)
+        ):
+            updated["issue_type"] = "onboarding"
+            issue_type = "onboarding"
+            reasons.append("override_issue_type=onboarding(onboarding_access_blocker)")
+
+    if issue_type is not None:
+        if "assignment_group" in allowed_fields:
+            desired_group = heuristic_assignment_group(text, issue_type)
+            if updated.get("assignment_group") != desired_group:
+                updated["assignment_group"] = desired_group
+                reasons.append(f"override_assignment_group={desired_group}")
+        if "resolution_action" in allowed_fields:
+            desired_resolution = heuristic_resolution_action(text, issue_type)
+            if updated.get("resolution_action") != desired_resolution:
+                updated["resolution_action"] = desired_resolution
+                reasons.append(f"override_resolution_action={desired_resolution}")
+
+    if "priority" in allowed_fields and updated.get("priority") is not None:
+        priority = updated["priority"]
+        has_critical_signal = any(keyword in text for keyword in CRITICAL_INCIDENT_KEYWORDS)
+        has_high_signal = any(keyword in text for keyword in HIGH_PRIORITY_SIGNAL_KEYWORDS)
+        if priority == "critical" and not has_critical_signal:
+            updated["priority"] = "high" if has_high_signal else "medium"
+            reasons.append(f"override_priority={updated['priority']}(deescalated_from_critical)")
+        elif (
+            priority == "high"
+            and issue_type in {"service_request", "onboarding"}
+            and not any(keyword in text for keyword in TIME_SENSITIVE_PRIORITY_KEYWORDS)
+        ):
+            updated["priority"] = "medium"
+            reasons.append("override_priority=medium(nonurgent_workflow_request)")
+        elif (
+            priority == "medium"
+            and issue_type == "identity_access"
+            and any(keyword in text for keyword in ("cannot sign in", "can't sign in", "2fa", "mfa", "locked"))
+            and not has_critical_signal
+        ):
+            updated["priority"] = "high"
+            reasons.append("override_priority=high(identity_lockout)")
+
+    return updated, reasons
+
+
 def build_action(
     ticket: dict, allowed_fields: list[str], instructions: str
 ) -> tuple[HelpdeskTicketAction, str, str | None]:
@@ -370,13 +641,50 @@ def build_action(
 
     try:
         llm_dict = call_llm(ticket, allowed_fields, instructions)
-        candidate = {
-            field: llm_dict[field]
-            for field in allowed_fields
-            if llm_dict.get(field) is not None
-        }
-        if not candidate:
+        validated_llm_fields: dict[str, Any] = {}
+        rejected_fields: list[str] = []
+        for field in allowed_fields:
+            value = llm_dict.get(field)
+            if value is None:
+                continue
+            try:
+                HelpdeskTicketAction(**{field: value})
+            except Exception:
+                rejected_fields.append(field)
+                continue
+            validated_llm_fields[field] = value
+
+        if not validated_llm_fields:
             raise ValueError("LLM returned no allowed fields")
+
+        candidate = heuristic_action(
+            ticket,
+            allowed_fields,
+            issue_type_override=validated_llm_fields.get("issue_type"),
+        )
+        candidate.update(validated_llm_fields)
+        accepted_fields = list(validated_llm_fields)
+        candidate, override_reasons = apply_domain_overrides(
+            ticket,
+            candidate,
+            allowed_fields,
+        )
+
+        backfilled_fields = [field for field in allowed_fields if field not in accepted_fields]
+        if backfilled_fields or rejected_fields or override_reasons:
+            reason_parts = []
+            if backfilled_fields:
+                reason_parts.append(f"heuristic_backfill={backfilled_fields}")
+            if rejected_fields:
+                reason_parts.append(f"invalid_llm_fields={rejected_fields}")
+            if override_reasons:
+                reason_parts.append(f"domain_overrides={override_reasons}")
+            return (
+                HelpdeskTicketAction(**candidate),
+                "llm_backfilled",
+                "; ".join(reason_parts),
+            )
+
         return HelpdeskTicketAction(**candidate), "llm", None
     except Exception as exc:
         return (
@@ -389,6 +697,10 @@ def build_action(
 def should_investigate(ticket: dict, history: list[dict[str, Any]]) -> tuple[bool, str | None]:
     if not ticket:
         return False, None
+    context_status = ticket.get("context_status") or {}
+    remaining_tools = context_status.get("remaining_tools") or []
+    if remaining_tools:
+        return True, str(remaining_tools[0])
     current_ticket_id = ticket.get("ticket_id")
     already_investigated = any(
         entry.get("ticket_id") == current_ticket_id
@@ -408,6 +720,22 @@ def merge_ticket_context(ticket: dict, observation: Any) -> dict:
     merged_ticket = dict(ticket)
     if getattr(observation, "last_tool_result", None) is not None:
         merged_ticket["last_tool_result"] = observation.last_tool_result
+    merged_ticket["recent_history"] = list(getattr(observation, "history", []))
+    merged_ticket["queue_position"] = getattr(observation, "queue_position", None)
+    merged_ticket["tickets_remaining"] = getattr(observation, "tickets_remaining", None)
+    merged_ticket["investigation_budget_remaining"] = getattr(
+        observation,
+        "investigation_budget_remaining",
+        None,
+    )
+    merged_ticket["average_score_so_far"] = getattr(observation, "average_score_so_far", None)
+    merged_ticket["progress_fraction"] = getattr(observation, "progress_fraction", None)
+    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"]
     return merged_ticket
 
 
@@ -518,7 +846,12 @@ def run() -> None:
                     ticket_id=ticket["ticket_id"],
                 )
 
-        final_reward = task_step_rewards[-1] if task_step_rewards else 0.0
+        final_rubric_reward = getattr(obs, "rubric_reward", None)
+        final_reward = (
+            float(final_rubric_reward)
+            if final_rubric_reward is not None
+            else (task_step_rewards[-1] if task_step_rewards else 0.0)
+        )
         all_results[task_id] = {
             "final_reward": final_reward,
             "step_count": step_num,

models.py

@@ -18,6 +18,7 @@ ASSIGNMENT_GROUP_SET = set(ASSIGNMENT_GROUPS)
 RESOLUTION_ACTION_SET = set(RESOLUTION_ACTIONS)
 ACTION_TYPE_SET = {"submit", "investigate"}
 TOOL_NAME_SET = {"lookup_related_ticket", "lookup_requester_history"}
+TOOL_NAME_SET.add("lookup_internal_routing_note")
 
 
 def _validate_choice(value: str, allowed: set[str], field_name: str) -> str:
@@ -113,6 +114,7 @@ class HelpdeskTicketObservation(Observation):
     task_name: str = ""
     instructions: str = ""
     allowed_fields: list[str] = Field(default_factory=list)
+    available_action_types: list[str] = Field(default_factory=list)
     available_tools: list[str] = Field(default_factory=list)
     investigation_budget_remaining: int = 0
     last_tool_result: Optional[dict[str, Any]] = None
@@ -122,7 +124,11 @@ class HelpdeskTicketObservation(Observation):
     tickets_after_current: int = 0
     tickets_processed: int = 0
     queue_position: int = 0
+    average_score_so_far: float = 0.0
+    progress_fraction: float = 0.0
     history: list[dict[str, Any]] = Field(default_factory=list)
+    last_reward_components: dict[str, Any] = Field(default_factory=dict)
+    rubric_reward: Optional[float] = None
 
 
 class HelpdeskTicketState(State):
@@ -136,7 +142,11 @@ class HelpdeskTicketState(State):
     # `reward` is the field the evaluator checks on GET /state (mentor spec)
     reward: Optional[float] = None
     done: bool = False
+    average_score_so_far: float = 0.0
     investigation_steps: int = 0
     investigation_budget_remaining: int = 0
+    investigation_penalty_applied: float = 0.0
     last_tool_result: Optional[dict[str, Any]] = None
+    last_reward_components: dict[str, Any] = Field(default_factory=dict)
+    ticket_tool_usage: dict[str, list[str]] = Field(default_factory=dict)
     history_entries: list[dict] = Field(default_factory=list)

policy_learning.py

@@ -0,0 +1,723 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+from dataclasses import asdict, dataclass
+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",
+)
+DEFAULT_SEARCH_POLICIES = (
+    "no_investigation",
+    "legacy_single_probe",
+    "investigate_when_context_hidden",
+    "context_chain",
+    "hybrid_context",
+)
+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
+
+
+POLICY_LIBRARY: dict[str, PolicyConfig] = {
+    "no_investigation": PolicyConfig(
+        name="no_investigation",
+        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",
+        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",
+        investigate_hidden_context=True,
+        investigate_related_ticket_hint=False,
+        investigate_ambiguity_history=False,
+        max_investigations_per_ticket=1,
+        description="Investigate once when the environment says context is hidden.",
+    ),
+    "context_chain": PolicyConfig(
+        name="context_chain",
+        investigate_hidden_context=True,
+        investigate_related_ticket_hint=False,
+        investigate_ambiguity_history=False,
+        max_investigations_per_ticket=3,
+        description="Follow the environment's required-tool chain until context is revealed.",
+    ),
+    "hybrid_context": PolicyConfig(
+        name="hybrid_context",
+        investigate_hidden_context=True,
+        investigate_related_ticket_hint=True,
+        investigate_ambiguity_history=True,
+        max_investigations_per_ticket=3,
+        description="Use hidden-context signals first, then legacy ambiguity hints.",
+    ),
+}
+
+
+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)
+    return HelpdeskTicketAction(**candidate)
+
+
+def choose_policy_action(
+    policy: PolicyConfig,
+    observation: HelpdeskTicketObservation,
+    investigations_by_ticket: dict[str, int],
+    submit_builder: SubmitBuilder,
+) -> tuple[HelpdeskTicketAction, str]:
+    ticket = observation.current_ticket or {}
+    ticket_id = str(ticket.get("ticket_id", ""))
+    ticket_investigations = investigations_by_ticket.get(ticket_id, 0)
+    revealed_tools = set(((ticket.get("context_status") or {}).get("revealed_tools") or []))
+    remaining_tools = list(((ticket.get("context_status") or {}).get("remaining_tools") or []))
+
+    if ticket_investigations < policy.max_investigations_per_ticket:
+        if policy.investigate_hidden_context and remaining_tools:
+            tool_name = str(remaining_tools[0])
+            return (
+                HelpdeskTicketAction(action_type="investigate", tool_name=tool_name),
+                "investigate_hidden_context",
+            )
+        if (
+            policy.investigate_related_ticket_hint
+            and ticket.get("related_ticket_id")
+            and "lookup_related_ticket" not in revealed_tools
+        ):
+            return (
+                HelpdeskTicketAction(
+                    action_type="investigate",
+                    tool_name="lookup_related_ticket",
+                ),
+                "investigate_related_ticket_hint",
+            )
+        if (
+            policy.investigate_ambiguity_history
+            and ticket.get("ambiguity_note")
+            and "lookup_requester_history" not in revealed_tools
+        ):
+            return (
+                HelpdeskTicketAction(
+                    action_type="investigate",
+                    tool_name="lookup_requester_history",
+                ),
+                "investigate_ambiguity_history",
+            )
+
+    return submit_builder(ticket, list(observation.allowed_fields)), "submit"
+
+
+def rollout_episode(
+    *,
+    env: HelpdeskTicketRoutingEnvironment,
+    policy: PolicyConfig,
+    seed: int,
+    task_id: int,
+    submit_builder: SubmitBuilder,
+) -> 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] = {}
+    episode_return = 0.0
+    trajectories: list[dict[str, Any]] = []
+
+    while not observation.done:
+        ticket = observation.current_ticket or {}
+        ticket_id = str(ticket.get("ticket_id", ""))
+        action, action_source = choose_policy_action(
+            policy,
+            observation,
+            investigations_by_ticket,
+            submit_builder,
+        )
+        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
+
+        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": 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,
+        "per_ticket_scores": list(env.state.per_ticket_scores),
+    }
+    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_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_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,
+) -> 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,
+            )
+            episode_summaries.append(summary)
+            trajectories.extend(episode_trajectories)
+
+    return {
+        "policy": policy.name,
+        "summary": summarize_policy_episodes(policy, episode_summaries),
+        "episodes": episode_summaries,
+        "trajectories": trajectories,
+    }
+
+
+def _selection_tuple(summary: dict[str, Any]) -> tuple[float, float, float, float]:
+    return (
+        float(summary["avg_normalized_return"]),
+        float(summary["avg_terminal_reward"]),
+        float(summary["avg_terminal_rubric_reward"]),
+        -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 = [
+        evaluate_policy(
+            policy,
+            seeds,
+            task_ids,
+            env_factory=env_factory,
+            submit_builder=submit_builder,
+        )
+        for policy in policies
+    ]
+    best_run = select_best_policy(policy_runs)
+    baseline_run = policy_runs[0]
+
+    report = {
+        "mode": "compare",
+        "task_ids": task_ids,
+        "seeds": seeds,
+        "selection_metric": "avg_normalized_return",
+        "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",
+            ),
+        },
+        "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,
+            )
+        ],
+        "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 = [
+        evaluate_policy(
+            policy,
+            train_seeds,
+            task_ids,
+            env_factory=env_factory,
+            submit_builder=submit_builder,
+        )
+        for policy in candidate_policies
+    ]
+    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,
+    )
+
+    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,
+    )
+
+    report = {
+        "mode": "search",
+        "task_ids": task_ids,
+        "train_seeds": train_seeds,
+        "eval_seeds": eval_seeds,
+        "selection_metric": "avg_normalized_return",
+        "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],
+        "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",
+            ),
+        },
+        "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_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()

pyproject.toml

@@ -24,12 +24,13 @@ dependencies = [
 
 [project.scripts]
 server = "server.app:main"
+policy-learn = "policy_learning:main"
 
 [project.optional-dependencies]
 dev = ["pytest", "httpx"]
 
 [tool.setuptools]
-py-modules = ["models", "client", "vocabulary"]
+py-modules = ["models", "client", "policy_learning", "vocabulary"]
 
 [tool.setuptools.packages.find]
 include = ["server*"]

server/environment.py

@@ -18,10 +18,68 @@ from server.tasks import get_task_definition, load_dataset
 
 
 QUEUE_SIZE_RANGE = (3, 5)
-AVAILABLE_TOOLS = ("lookup_related_ticket", "lookup_requester_history")
+AVAILABLE_ACTION_TYPES = ("submit", "investigate")
+AVAILABLE_TOOLS = (
+    "lookup_related_ticket",
+    "lookup_requester_history",
+    "lookup_internal_routing_note",
+)
 FREE_INVESTIGATIONS_PER_TICKET = 1
 EXTRA_INVESTIGATION_COST = 0.02
 MAX_EXTRA_INVESTIGATION_PENALTY = 0.15
+USEFUL_INVESTIGATION_REWARD = 0.08
+PREMATURE_SUBMIT_PENALTY = 0.10
+
+TASK3_INVESTIGATION_TOOL_PLAN: dict[str, tuple[str, ...]] = {
+    "ticket-021": ("lookup_related_ticket", "lookup_requester_history"),
+    "ticket-022": ("lookup_internal_routing_note",),
+    "ticket-027": ("lookup_internal_routing_note",),
+    "ticket-029": ("lookup_internal_routing_note",),
+    "ticket-038": ("lookup_related_ticket", "lookup_requester_history"),
+    "ticket-045": ("lookup_related_ticket", "lookup_requester_history"),
+    "TKT-NONDEFAULT-001": ("lookup_internal_routing_note",),
+    "TKT-NONDEFAULT-002": ("lookup_internal_routing_note",),
+    "TKT-NONDEFAULT-003": ("lookup_internal_routing_note",),
+}
+
+HARD_TASK_DESCRIPTION_REDACTIONS: dict[str, str] = {
+    "ticket-021": (
+        "Production checkout is still unstable after a recent fix. "
+        "Additional routing context is available via investigation."
+    ),
+    "ticket-022": (
+        "Usage charges increased while the integration was failing. "
+        "Additional routing context is available via investigation."
+    ),
+    "ticket-027": (
+        "A vendor offer arrived with a near-term deadline. "
+        "Additional routing context is available via investigation."
+    ),
+    "ticket-029": (
+        "A team needs a large seat expansion right away. "
+        "Additional routing context is available via investigation."
+    ),
+    "ticket-038": (
+        "A prior invoice discrepancy is still unresolved and now time-sensitive. "
+        "Additional routing context is available via investigation."
+    ),
+    "ticket-045": (
+        "A company-wide suspension remains unresolved after repeated follow-ups. "
+        "Additional routing context is available via investigation."
+    ),
+    "TKT-NONDEFAULT-001": (
+        "A user needs help with a billing-style question. "
+        "Additional routing context is available via investigation."
+    ),
+    "TKT-NONDEFAULT-002": (
+        "A client compliance scan surfaced a product-specific issue. "
+        "Additional routing context is available via investigation."
+    ),
+    "TKT-NONDEFAULT-003": (
+        "A contractor onboarding workflow is blocked by an account problem. "
+        "Additional routing context is available via investigation."
+    ),
+}
 
 
 def _coerce_optional_int(value: Any, field_name: str) -> Optional[int]:
@@ -86,7 +144,11 @@ class HelpdeskTicketRoutingEnvironment(
             current_ticket_index=0,
             per_ticket_scores=[],
             total_reward=0.0,
+            average_score_so_far=0.0,
             investigation_budget_remaining=queue_size * FREE_INVESTIGATIONS_PER_TICKET,
+            investigation_penalty_applied=0.0,
+            last_reward_components={},
+            ticket_tool_usage={},
         )
 
         return self._build_observation(task)
@@ -122,54 +184,104 @@ class HelpdeskTicketRoutingEnvironment(
         if extra_fields:
             # Penalty: record score 0.0, advance index, return penalty observation
             self._state.per_ticket_scores.append(0.0)
-            self._state.history_entries.append(
-                self._build_history_entry(
-                    current_ticket,
-                    predicted=action.model_dump(exclude_none=True),
-                    score=0.0,
-                    breakdown={},
-                    queue_position=idx + 1,
-                    penalty_reason=f"extra_fields: {sorted(extra_fields)}",
-                )
-            )
+            self._state.average_score_so_far = self._current_average_score()
             self._state.step_count += 1
             self._state.current_ticket_index += 1
             is_done = self._state.current_ticket_index >= len(self._queue)
             self._state.done = is_done
+            trajectory_reward = None
+            investigation_penalty = self._compute_episode_penalty() if is_done else 0.0
             if is_done:
-                traj_reward = compute_trajectory_reward(
+                trajectory_reward = compute_trajectory_reward(
                     self._state.per_ticket_scores, len(self._queue), self._state.step_count
                 )
-                final_reward = self._apply_episode_economics(traj_reward)
+                final_reward = self._apply_episode_economics(trajectory_reward)
                 self._state.total_reward = final_reward
             else:
                 final_reward = 0.0
+            reward_components = self._build_reward_components(
+                ticket_score=0.0,
+                field_breakdown={},
+                shaped_step_reward=0.0,
+                reward_kind="trajectory" if is_done else "step_penalty",
+                final_reward=final_reward,
+                trajectory_reward=trajectory_reward,
+                investigation_penalty=investigation_penalty,
+                penalty_reason=f"extra_fields: {sorted(extra_fields)}",
+            )
+            self._state.history_entries.append(
+                self._build_history_entry(
+                    current_ticket,
+                    predicted=action.model_dump(exclude_none=True),
+                    score=0.0,
+                    breakdown={},
+                    queue_position=idx + 1,
+                    reward=final_reward,
+                    rubric_reward=final_reward if is_done else None,
+                    reward_kind="trajectory" if is_done else "step_penalty",
+                    penalty_reason=f"extra_fields: {sorted(extra_fields)}",
+                    reward_components=reward_components,
+                )
+            )
             self._state.last_step_reward = final_reward
             self._state.reward = final_reward
+            self._state.investigation_penalty_applied = self._compute_episode_penalty()
             self._state.last_tool_result = None
-            return self._build_observation(task, done=is_done, reward=final_reward)
+            self._state.last_reward_components = reward_components
+            return self._build_observation(
+                task,
+                done=is_done,
+                reward=final_reward,
+                rubric_reward=final_reward if is_done else None,
+            )
 
         score, breakdown = grade_action(action, current_ticket, task_id)
         step_reward = compute_step_reward(score)
+        context_penalty, missing_required_tools = self._submit_context_penalty(current_ticket)
+        milestone_adjustment = step_reward - score
 
         is_done = (self._state.current_ticket_index + 1) >= len(self._queue)
+        trajectory_reward = None
+        investigation_penalty = 0.0
+        rubric_reward = None
 
         if is_done:
             self._state.per_ticket_scores.append(score)
+            self._state.average_score_so_far = self._current_average_score()
             self._state.step_count += 1
             self._state.current_ticket_index += 1
-            traj_reward = compute_trajectory_reward(
+            trajectory_reward = compute_trajectory_reward(
                 self._state.per_ticket_scores,
                 len(self._queue),
                 self._state.step_count,
             )
-            final_reward = self._apply_episode_economics(traj_reward)
-            self._state.total_reward = final_reward
+            rubric_reward = self._apply_episode_economics(trajectory_reward)
+            final_reward = max(0.0, min(1.0, rubric_reward - context_penalty))
+            self._state.total_reward = rubric_reward
+            investigation_penalty = self._compute_episode_penalty()
         else:
             self._state.per_ticket_scores.append(score)
+            self._state.average_score_so_far = self._current_average_score()
             self._state.step_count += 1
             self._state.current_ticket_index += 1
-            final_reward = step_reward
+            final_reward = max(0.0, min(1.0, step_reward - context_penalty))
+
+        reward_components = self._build_reward_components(
+            ticket_score=score,
+            field_breakdown=breakdown,
+            shaped_step_reward=step_reward,
+            reward_kind="trajectory" if is_done else "step",
+            final_reward=final_reward,
+            milestone_adjustment=milestone_adjustment,
+            trajectory_reward=trajectory_reward,
+            investigation_penalty=investigation_penalty,
+            extra_details={
+                "context_gap_penalty": context_penalty,
+                "required_tools": self._required_tools_for_ticket(current_ticket),
+                "remaining_required_tools": missing_required_tools,
+                "rubric_reward": rubric_reward,
+            },
+        )
 
         history_entry = self._build_history_entry(
             current_ticket,
@@ -177,15 +289,26 @@ class HelpdeskTicketRoutingEnvironment(
             score=score,
             breakdown=breakdown,
             queue_position=idx + 1,
+            reward=final_reward,
+            rubric_reward=rubric_reward if is_done else None,
+            reward_kind="trajectory" if is_done else "step",
+            reward_components=reward_components,
         )
         self._state.history_entries.append(history_entry)
 
         self._state.last_step_reward = final_reward
         self._state.reward = final_reward
         self._state.done = is_done
+        self._state.investigation_penalty_applied = self._compute_episode_penalty()
         self._state.last_tool_result = None
+        self._state.last_reward_components = reward_components
 
-        return self._build_observation(task, done=is_done, reward=final_reward)
+        return self._build_observation(
+            task,
+            done=is_done,
+            reward=final_reward,
+            rubric_reward=rubric_reward if is_done else None,
+        )
 
     @property
     def state(self) -> HelpdeskTicketState:
@@ -195,15 +318,112 @@ class HelpdeskTicketRoutingEnvironment(
     # Helpers
     # ------------------------------------------------------------------
 
-    def _apply_episode_economics(self, base_reward: float) -> float:
+    def _compute_episode_penalty(self) -> float:
         free_investigations = len(self._queue) * FREE_INVESTIGATIONS_PER_TICKET
         extra_investigations = max(0, self._state.investigation_steps - free_investigations)
-        penalty = min(
+        return min(
             MAX_EXTRA_INVESTIGATION_PENALTY,
             extra_investigations * EXTRA_INVESTIGATION_COST,
         )
+
+    def _apply_episode_economics(self, base_reward: float) -> float:
+        penalty = self._compute_episode_penalty()
         return max(0.0, min(1.0, base_reward - penalty))
 
+    def _current_average_score(self) -> float:
+        if not self._state.per_ticket_scores:
+            return 0.0
+        return sum(self._state.per_ticket_scores) / len(self._state.per_ticket_scores)
+
+    def _required_tools_for_ticket(
+        self,
+        ticket: HelpdeskTicketRecord,
+        task_id: int | None = None,
+    ) -> list[str]:
+        resolved_task_id = self._state.current_task_id if task_id is None else task_id
+        if resolved_task_id != 3:
+            return []
+        return list(TASK3_INVESTIGATION_TOOL_PLAN.get(ticket.ticket_id, ()))
+
+    def _used_tools_for_ticket(self, ticket_id: str) -> list[str]:
+        return list(self._state.ticket_tool_usage.get(ticket_id, []))
+
+    def _remaining_tools_for_ticket(
+        self,
+        ticket: HelpdeskTicketRecord,
+        task_id: int | None = None,
+    ) -> list[str]:
+        required_tools = self._required_tools_for_ticket(ticket, task_id)
+        used_tools = set(self._used_tools_for_ticket(ticket.ticket_id))
+        return [tool for tool in required_tools if tool not in used_tools]
+
+    def _record_tool_usage(self, ticket_id: str, tool_name: str) -> None:
+        used = self._state.ticket_tool_usage.setdefault(ticket_id, [])
+        if tool_name not in used:
+            used.append(tool_name)
+
+    def _investigation_hints_for_ticket(self, ticket: HelpdeskTicketRecord) -> list[str]:
+        hints: list[str] = []
+        remaining_tools = self._remaining_tools_for_ticket(ticket)
+        if "lookup_internal_routing_note" in remaining_tools:
+            hints.append("An internal routing note may disambiguate the correct workflow.")
+        if "lookup_related_ticket" in remaining_tools:
+            hints.append("A linked prior ticket can reveal important follow-up context.")
+        if "lookup_requester_history" in remaining_tools:
+            hints.append("Requester history may clarify severity or routing intent.")
+        return hints
+
+    def _visible_description(self, ticket: HelpdeskTicketRecord) -> str:
+        if (
+            self._state.current_task_id == 3
+            and self._remaining_tools_for_ticket(ticket)
+            and ticket.ticket_id in HARD_TASK_DESCRIPTION_REDACTIONS
+        ):
+            return HARD_TASK_DESCRIPTION_REDACTIONS[ticket.ticket_id]
+        return ticket.description
+
+    def _submit_context_penalty(self, ticket: HelpdeskTicketRecord) -> tuple[float, list[str]]:
+        required_tools = self._required_tools_for_ticket(ticket)
+        if not required_tools:
+            return 0.0, []
+        remaining_tools = self._remaining_tools_for_ticket(ticket)
+        if not remaining_tools:
+            return 0.0, []
+        penalty = PREMATURE_SUBMIT_PENALTY * (len(remaining_tools) / len(required_tools))
+        return penalty, remaining_tools
+
+    def _build_reward_components(
+        self,
+        *,
+        ticket_score: float,
+        field_breakdown: dict[str, float],
+        shaped_step_reward: float,
+        reward_kind: str,
+        final_reward: float,
+        milestone_adjustment: float = 0.0,
+        trajectory_reward: float | None = None,
+        investigation_penalty: float = 0.0,
+        penalty_reason: str | None = None,
+        extra_details: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        components: dict[str, Any] = {
+            "reward_kind": reward_kind,
+            "ticket_score": ticket_score,
+            "field_breakdown": field_breakdown,
+            "shaped_step_reward": shaped_step_reward,
+            "milestone_adjustment": milestone_adjustment,
+            "final_reward": final_reward,
+            "average_score_so_far": self._current_average_score(),
+            "investigation_penalty_applied": investigation_penalty,
+        }
+        if trajectory_reward is not None:
+            components["trajectory_reward"] = trajectory_reward
+        if penalty_reason is not None:
+            components["penalty_reason"] = penalty_reason
+        if extra_details:
+            components.update(extra_details)
+        return components
+
     def _lookup_related_ticket(
         self,
         current_ticket: HelpdeskTicketRecord,
@@ -259,6 +479,15 @@ class HelpdeskTicketRoutingEnvironment(
             "matches": matches,
         }
 
+    def _lookup_internal_routing_note(self, current_ticket: HelpdeskTicketRecord) -> dict[str, Any]:
+        found = current_ticket.ambiguity_note is not None
+        return {
+            "tool_name": "lookup_internal_routing_note",
+            "found": found,
+            "ticket_id": current_ticket.ticket_id,
+            "routing_note": current_ticket.ambiguity_note if found else "",
+        }
+
     def _run_investigation_tool(
         self,
         current_ticket: HelpdeskTicketRecord,
@@ -269,6 +498,8 @@ class HelpdeskTicketRoutingEnvironment(
             return self._lookup_related_ticket(current_ticket, target_ticket_id)
         if tool_name == "lookup_requester_history":
             return self._lookup_requester_history(current_ticket)
+        if tool_name == "lookup_internal_routing_note":
+            return self._lookup_internal_routing_note(current_ticket)
         raise ValueError(f"Unsupported tool_name: {tool_name}")
 
     def _handle_investigation_action(
@@ -296,6 +527,14 @@ class HelpdeskTicketRoutingEnvironment(
             action.tool_name,
             action.tool_target_ticket_id,
         )
+        required_tools = self._required_tools_for_ticket(current_ticket)
+        already_used = action.tool_name in self._used_tools_for_ticket(current_ticket.ticket_id)
+        useful_investigation = (
+            action.tool_name in required_tools
+            and not already_used
+            and bool(tool_result.get("found", True))
+        )
+        self._record_tool_usage(current_ticket.ticket_id, action.tool_name)
         self._state.step_count += 1
         self._state.investigation_steps += 1
         self._state.investigation_budget_remaining = max(
@@ -303,9 +542,25 @@ class HelpdeskTicketRoutingEnvironment(
             self._state.investigation_budget_remaining - 1,
         )
         self._state.last_tool_result = tool_result
-        self._state.last_step_reward = 0.0
-        self._state.reward = 0.0
+        investigation_reward = USEFUL_INVESTIGATION_REWARD if useful_investigation else 0.0
+        self._state.last_step_reward = investigation_reward
+        self._state.reward = investigation_reward
         self._state.done = False
+        self._state.investigation_penalty_applied = self._compute_episode_penalty()
+        reward_components = self._build_reward_components(
+            ticket_score=0.0,
+            field_breakdown={},
+            shaped_step_reward=investigation_reward,
+            reward_kind="investigation",
+            final_reward=investigation_reward,
+            investigation_penalty=self._state.investigation_penalty_applied,
+            extra_details={
+                "new_context_revealed": useful_investigation,
+                "required_tools": required_tools,
+                "remaining_required_tools": self._remaining_tools_for_ticket(current_ticket),
+                "tool_name": action.tool_name,
+            },
+        )
         self._state.history_entries.append(
             self._build_history_entry(
                 current_ticket,
@@ -313,21 +568,35 @@ class HelpdeskTicketRoutingEnvironment(
                 score=0.0,
                 breakdown={},
                 queue_position=idx + 1,
+                reward=investigation_reward,
+                reward_kind="investigation",
                 tool_result=tool_result,
+                reward_components=reward_components,
             )
         )
-        return self._build_observation(task, done=False, reward=0.0)
+        self._state.last_reward_components = reward_components
+        return self._build_observation(task, done=False, reward=investigation_reward)
 
     def _build_ticket_view(self, ticket: HelpdeskTicketRecord) -> dict[str, Any]:
+        required_tools = self._required_tools_for_ticket(ticket)
+        revealed_tools = self._used_tools_for_ticket(ticket.ticket_id)
+        remaining_tools = self._remaining_tools_for_ticket(ticket)
         ticket_view: dict[str, Any] = {
             "ticket_id": ticket.ticket_id,
             "title": ticket.title,
             "requester": ticket.requester,
-            "description": ticket.description,
+            "description": self._visible_description(ticket),
         }
-        if ticket.ambiguity_note is not None:
+        if required_tools:
+            ticket_view["context_status"] = {
+                "investigation_required": True,
+                "revealed_tools": revealed_tools,
+                "remaining_tools": remaining_tools,
+                "hints": self._investigation_hints_for_ticket(ticket),
+            }
+        if ticket.ambiguity_note is not None and "lookup_internal_routing_note" not in remaining_tools:
             ticket_view["ambiguity_note"] = ticket.ambiguity_note
-        if ticket.related_ticket_id is not None:
+        if ticket.related_ticket_id is not None and "lookup_related_ticket" not in remaining_tools:
             ticket_view["related_ticket_id"] = ticket.related_ticket_id
             related_ticket = self._tickets_by_id.get(ticket.related_ticket_id)
             if related_ticket is not None:
@@ -339,6 +608,50 @@ class HelpdeskTicketRoutingEnvironment(
                 }
         return ticket_view
 
+    def _build_feedback_summary(
+        self,
+        *,
+        predicted: dict[str, Any],
+        score: float,
+        breakdown: dict[str, float],
+        reward: float | None = None,
+        rubric_reward: float | None = None,
+        reward_kind: str | None = None,
+        penalty_reason: str | None = None,
+        tool_result: dict[str, Any] | None = None,
+        reward_components: dict[str, Any] | None = None,
+    ) -> str:
+        parts: list[str] = []
+
+        if reward_kind == "investigation":
+            tool_name = predicted.get("tool_name") or (tool_result or {}).get("tool_name")
+            parts.append(f"Investigation step used {tool_name or 'a tool'}")
+            if reward_components and reward_components.get("new_context_revealed"):
+                parts.append("new context was revealed")
+        elif penalty_reason is not None:
+            parts.append(f"Penalty applied: {penalty_reason}")
+        else:
+            parts.append(f"Ticket score={score:.2f}")
+
+        if breakdown:
+            field_scores = ", ".join(
+                f"{field}={value:.2f}" for field, value in sorted(breakdown.items())
+            )
+            parts.append(f"field_scores[{field_scores}]")
+        if reward is not None:
+            parts.append(f"reward={reward:.2f}")
+        if rubric_reward is not None:
+            parts.append(f"rubric_reward={rubric_reward:.2f}")
+        if reward_components:
+            context_gap_penalty = reward_components.get("context_gap_penalty")
+            if context_gap_penalty:
+                parts.append(f"context_gap_penalty={context_gap_penalty:.2f}")
+            remaining_required_tools = reward_components.get("remaining_required_tools") or []
+            if remaining_required_tools:
+                parts.append(f"missing_context={remaining_required_tools}")
+
+        return "; ".join(parts)
+
     def _build_history_entry(
         self,
         ticket: HelpdeskTicketRecord,
@@ -347,9 +660,15 @@ class HelpdeskTicketRoutingEnvironment(
         score: float,
         breakdown: dict[str, float],
         queue_position: int,
+        reward: float | None = None,
+        rubric_reward: float | None = None,
+        reward_kind: str | None = None,
         penalty_reason: str | None = None,
         tool_result: dict[str, Any] | None = None,
+        reward_components: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
+        remaining_tools = self._remaining_tools_for_ticket(ticket)
+        revealed_tools = self._used_tools_for_ticket(ticket.ticket_id)
         history_entry: dict[str, Any] = {
             "ticket_id": ticket.ticket_id,
             "title": ticket.title,
@@ -359,9 +678,15 @@ class HelpdeskTicketRoutingEnvironment(
             "breakdown": breakdown,
             "queue_position": queue_position,
         }
-        if ticket.ambiguity_note is not None:
+        if reward is not None:
+            history_entry["reward"] = reward
+        if rubric_reward is not None:
+            history_entry["rubric_reward"] = rubric_reward
+        if reward_kind is not None:
+            history_entry["reward_kind"] = reward_kind
+        if ticket.ambiguity_note is not None and "lookup_internal_routing_note" not in remaining_tools:
             history_entry["ambiguity_note"] = ticket.ambiguity_note
-        if ticket.related_ticket_id is not None:
+        if ticket.related_ticket_id is not None and "lookup_related_ticket" not in remaining_tools:
             history_entry["related_ticket_id"] = ticket.related_ticket_id
             related_ticket = self._tickets_by_id.get(ticket.related_ticket_id)
             if related_ticket is not None:
@@ -375,6 +700,21 @@ class HelpdeskTicketRoutingEnvironment(
             history_entry["penalty_reason"] = penalty_reason
         if tool_result is not None:
             history_entry["tool_result"] = tool_result
+        if reward_components is not None:
+            history_entry["reward_components"] = reward_components
+        if revealed_tools:
+            history_entry["revealed_tools"] = revealed_tools
+        history_entry["feedback_summary"] = self._build_feedback_summary(
+            predicted=predicted,
+            score=score,
+            breakdown=breakdown,
+            reward=reward,
+            rubric_reward=rubric_reward,
+            reward_kind=reward_kind,
+            penalty_reason=penalty_reason,
+            tool_result=tool_result,
+            reward_components=reward_components,
+        )
         return history_entry
 
     def _build_observation(
@@ -382,6 +722,7 @@ class HelpdeskTicketRoutingEnvironment(
         task: dict,
         done: bool = False,
         reward: float | None = None,
+        rubric_reward: float | None = None,
     ) -> HelpdeskTicketObservation:
         idx = self._state.current_ticket_index
         queue_size = len(self._queue)
@@ -395,28 +736,47 @@ class HelpdeskTicketRoutingEnvironment(
             queue_position = 0
 
         history = list(self._state.history_entries)
+        last_history_entry = history[-1] if history else None
         tickets_remaining = max(0, queue_size - idx)
         tickets_after_current = max(
             0,
             tickets_remaining - (1 if ticket_view is not None else 0),
         )
+        progress_fraction = (idx / queue_size) if queue_size else 0.0
+
+        metadata = {
+            "queue_position": queue_position,
+            "tickets_remaining_includes_current": ticket_view is not None,
+            "has_ambiguity_note": bool(ticket_view and ticket_view.get("ambiguity_note")),
+            "has_related_ticket_context": bool(
+                ticket_view and ticket_view.get("related_ticket_preview")
+            ),
+            "action_mode": "investigate_or_submit",
+            "available_action_types": list(AVAILABLE_ACTION_TYPES),
+            "average_score_so_far": self._state.average_score_so_far,
+            "progress_fraction": progress_fraction,
+            "investigation_penalty_applied": self._state.investigation_penalty_applied,
+        }
+        if last_history_entry is not None:
+            metadata["last_score"] = last_history_entry.get("score")
+            metadata["last_reward"] = last_history_entry.get("reward")
+            metadata["last_reward_kind"] = last_history_entry.get("reward_kind")
+            metadata["last_breakdown"] = last_history_entry.get("breakdown")
+            metadata["last_feedback_summary"] = last_history_entry.get("feedback_summary")
+            metadata["last_reward_components"] = last_history_entry.get("reward_components", {})
+            if "penalty_reason" in last_history_entry:
+                metadata["last_penalty_reason"] = last_history_entry["penalty_reason"]
 
         return HelpdeskTicketObservation(
             done=done,
             reward=reward,
-            metadata={
-                "queue_position": queue_position,
-                "tickets_remaining_includes_current": ticket_view is not None,
-                "has_ambiguity_note": bool(ticket_view and ticket_view.get("ambiguity_note")),
-                "has_related_ticket_context": bool(
-                    ticket_view and ticket_view.get("related_ticket_preview")
-                ),
-                "action_mode": "investigate_or_submit",
-            },
+            rubric_reward=rubric_reward,
+            metadata=metadata,
             task_id=task["id"],
             task_name=task["name"],
             instructions=task["instructions"],
             allowed_fields=list(task["allowed_fields"]),
+            available_action_types=list(AVAILABLE_ACTION_TYPES),
             available_tools=list(AVAILABLE_TOOLS),
             investigation_budget_remaining=self._state.investigation_budget_remaining,
             last_tool_result=self._state.last_tool_result,
@@ -426,5 +786,8 @@ class HelpdeskTicketRoutingEnvironment(
             tickets_after_current=tickets_after_current,
             tickets_processed=idx,
             queue_position=queue_position,
+            average_score_so_far=self._state.average_score_so_far,
+            progress_fraction=progress_fraction,
             history=history,
+            last_reward_components=dict(self._state.last_reward_components),
         )

server/tasks.py

@@ -37,7 +37,9 @@ TASKS = {
             "Perform full helpdesk routing by selecting the best issue type, "
             "priority, assignment group, and resolution action for the ticket. "
             "Use any ambiguity notes or related-ticket previews when present. "
-            "You may investigate with tools before you submit the final action."
+            "Some hard tickets intentionally hide decisive routing context until "
+            "you investigate with the available tools, so premature submission can "
+            "underperform even when the visible text looks plausible."
         ),
         "allowed_fields": [
             "issue_type",

tests/test_api_integration.py

@@ -167,6 +167,9 @@ class TestResetEndpoint(unittest.TestCase):
     def test_reset_reward_is_null(self):
         self.assertIsNone(self.data["reward"])
 
+    def test_reset_rubric_reward_is_null(self):
+        self.assertIsNone(self.data["rubric_reward"])
+
     def test_reset_task_id_is_1(self):
         self.assertEqual(self.data["task_id"], 1)
 
@@ -177,6 +180,13 @@ class TestResetEndpoint(unittest.TestCase):
         self.assertIsInstance(self.data["allowed_fields"], list)
         self.assertGreater(len(self.data["allowed_fields"]), 0)
 
+    def test_reset_available_action_types_exposed(self):
+        self.assertEqual(self.data["available_action_types"], ["submit", "investigate"])
+
+    def test_reset_progress_metrics_start_at_zero(self):
+        self.assertEqual(self.data["average_score_so_far"], 0.0)
+        self.assertEqual(self.data["progress_fraction"], 0.0)
+
 
 class TestStepEndpoint(unittest.TestCase):
     """2.1.4 — POST /step returns observation JSON with reward in [0.0, 1.0]."""
@@ -200,6 +210,35 @@ class TestStepEndpoint(unittest.TestCase):
     def test_step_tickets_processed_is_1(self):
         self.assertEqual(self.data["tickets_processed"], 1)
 
+    def test_step_metadata_exposes_last_feedback_summary(self):
+        metadata = self.data.get("metadata", {})
+        self.assertIn("last_feedback_summary", metadata)
+        self.assertIsInstance(metadata["last_feedback_summary"], str)
+        self.assertTrue(metadata["last_feedback_summary"])
+
+    def test_step_history_entry_includes_feedback_summary(self):
+        history = self.data.get("history", [])
+        self.assertGreater(len(history), 0)
+        self.assertIn("feedback_summary", history[-1])
+        self.assertIsInstance(history[-1]["feedback_summary"], str)
+        self.assertTrue(history[-1]["feedback_summary"])
+
+    def test_step_exposes_structured_reward_components(self):
+        self.assertIn("last_reward_components", self.data)
+        self.assertIsInstance(self.data["last_reward_components"], dict)
+        self.assertIn("ticket_score", self.data["last_reward_components"])
+        self.assertIn("final_reward", self.data["last_reward_components"])
+        self.assertEqual(
+            self.data["metadata"].get("last_reward_components"),
+            self.data["last_reward_components"],
+        )
+
+    def test_step_progress_metrics_are_exposed(self):
+        self.assertIn("average_score_so_far", self.data)
+        self.assertIn("progress_fraction", self.data)
+        self.assertGreaterEqual(self.data["progress_fraction"], 0.0)
+        self.assertLessEqual(self.data["progress_fraction"], 1.0)
+
 
 class TestStateEndpoint(unittest.TestCase):
     """2.1.5 — GET /state returns current episode state JSON after a reset."""
@@ -278,6 +317,38 @@ class TestFullSeededEpisode(unittest.TestCase):
         self.assertGreaterEqual(final_reward, 0.0)
         self.assertLessEqual(final_reward, 1.0)
 
+    def test_full_episode_terminal_rubric_reward_in_unit_interval(self):
+        reset_resp = _reset(task_id=1, seed=42)
+        self.assertEqual(reset_resp.status_code, 200)
+        obs = reset_resp.json()
+
+        allowed_fields = obs["allowed_fields"]
+        final_rubric_reward = None
+        for _ in range(20):
+            action_payload: dict = {}
+            if "issue_type" in allowed_fields:
+                action_payload["issue_type"] = "general_inquiry"
+            if "priority" in allowed_fields:
+                action_payload["priority"] = "medium"
+            if "assignment_group" in allowed_fields:
+                action_payload["assignment_group"] = "service_desk"
+            if "resolution_action" in allowed_fields:
+                action_payload["resolution_action"] = "acknowledge"
+
+            step_resp = client.post("/step", json=action_payload)
+            self.assertEqual(step_resp.status_code, 200)
+            obs = step_resp.json()
+
+            if obs["done"]:
+                final_rubric_reward = obs.get("rubric_reward")
+                break
+
+        self.assertIsNotNone(
+            final_rubric_reward, "Terminal observation did not include rubric_reward"
+        )
+        self.assertGreaterEqual(final_rubric_reward, 0.0)
+        self.assertLessEqual(final_rubric_reward, 1.0)
+
     def test_full_episode_all_tasks_complete(self):
         """4.1.1 — Full seeded episode completes for each task ID (1, 2, 3)."""
         for task_id in (1, 2, 3):

tests/test_competitive_upgrade.py

@@ -182,6 +182,16 @@ class TestStateHasRewardAndDone(unittest.TestCase):
             obs = env.step(_heuristic_action(obs))
             self.assertFalse(env.state.done)
 
+    def test_state_tracks_average_score_and_reward_components(self) -> None:
+        env = _make_env()
+        obs = env.reset(seed=42, task_id=1)
+        env.step(_heuristic_action(obs))
+        state = env.state
+        self.assertGreaterEqual(state.average_score_so_far, 0.0)
+        self.assertLessEqual(state.average_score_so_far, 1.0)
+        self.assertIsInstance(state.last_reward_components, dict)
+        self.assertIn("final_reward", state.last_reward_components)
+
 
 # ---------------------------------------------------------------------------
 # 9.3 — History entry contains title and predicted
@@ -318,7 +328,7 @@ class TestAmbiguityNoteInObservation(unittest.TestCase):
                 return seed
         return None
 
-    def test_ambiguity_note_present_when_ticket_has_one(self) -> None:
+    def test_ambiguity_note_hidden_until_internal_note_lookup(self) -> None:
         """Force a ticket with ambiguity_note by patching the dataset."""
         from unittest.mock import patch
         from server.tasks import load_dataset
@@ -336,8 +346,22 @@ class TestAmbiguityNoteInObservation(unittest.TestCase):
             obs = env.reset(seed=0, task_id=3)
 
         self.assertIsNotNone(obs.current_ticket)
-        self.assertIn("ambiguity_note", obs.current_ticket)
+        self.assertNotIn("ambiguity_note", obs.current_ticket)
+        self.assertIn("context_status", obs.current_ticket)
+        self.assertIn(
+            "lookup_internal_routing_note",
+            obs.current_ticket["context_status"]["remaining_tools"],
+        )
+
+        obs = env.step(
+            HelpdeskTicketAction(
+                action_type="investigate",
+                tool_name="lookup_internal_routing_note",
+            )
+        )
+
         self.assertEqual(obs.current_ticket["ambiguity_note"], target.ambiguity_note)
+        self.assertGreater(obs.reward or 0.0, 0.0)
 
     def test_ambiguity_note_absent_when_ticket_has_none(self) -> None:
         """Tickets without ambiguity_note should not expose the key."""
@@ -370,6 +394,13 @@ class TestAmbiguityNoteInObservation(unittest.TestCase):
         with patch.object(env, "_dataset", [ticket]):
             obs = env.reset(seed=0, task_id=3)
 
+        self.assertNotIn("ambiguity_note", obs.current_ticket)
+        obs = env.step(
+            HelpdeskTicketAction(
+                action_type="investigate",
+                tool_name="lookup_internal_routing_note",
+            )
+        )
         self.assertIn("ambiguity_note", obs.current_ticket)
 
 
@@ -397,12 +428,27 @@ class TestRelatedTicketPreviewInObservation(unittest.TestCase):
             ):
                 obs = env.reset(seed=0, task_id=3, queue_size=1)
 
-        return env, obs, related
+        return env, obs, ticket, related
 
     def test_related_ticket_preview_present_when_ticket_has_link(self) -> None:
-        env, obs, related = self._reset_linked_ticket_env()
+        env, obs, ticket, related = self._reset_linked_ticket_env()
 
         self.assertIsNotNone(obs.current_ticket)
+        self.assertNotIn("related_ticket_preview", obs.current_ticket)
+        self.assertIn("context_status", obs.current_ticket)
+        self.assertIn(
+            "lookup_related_ticket",
+            obs.current_ticket["context_status"]["remaining_tools"],
+        )
+
+        obs = env.step(
+            HelpdeskTicketAction(
+                action_type="investigate",
+                tool_name="lookup_related_ticket",
+                tool_target_ticket_id=ticket.related_ticket_id,
+            )
+        )
+
         self.assertIn("related_ticket_preview", obs.current_ticket)
         self.assertEqual(
             obs.current_ticket["related_ticket_preview"]["ticket_id"],
@@ -414,8 +460,22 @@ class TestRelatedTicketPreviewInObservation(unittest.TestCase):
         )
 
     def test_history_keeps_related_ticket_preview_after_step(self) -> None:
-        env, obs, related = self._reset_linked_ticket_env()
-        next_obs = env.step(_heuristic_action(obs))
+        env, obs, ticket, related = self._reset_linked_ticket_env()
+        env.step(
+            HelpdeskTicketAction(
+                action_type="investigate",
+                tool_name="lookup_related_ticket",
+                tool_target_ticket_id=ticket.related_ticket_id,
+            )
+        )
+        next_obs = env.step(
+            HelpdeskTicketAction(
+                issue_type=ticket.issue_type,
+                priority=ticket.priority,
+                assignment_group=ticket.assignment_group,
+                resolution_action=ticket.resolution_action,
+            )
+        )
 
         self.assertGreaterEqual(len(next_obs.history), 1)
         self.assertIn("related_ticket_preview", next_obs.history[0])
@@ -563,6 +623,58 @@ class TestInvestigationActions(unittest.TestCase):
         self.assertTrue(obs2.last_tool_result["found"])
         self.assertGreaterEqual(len(obs2.last_tool_result["matches"]), 1)
 
+    def test_internal_note_tool_reveals_hidden_hard_task_context(self) -> None:
+        from unittest.mock import patch
+
+        dataset = load_dataset()
+        ticket = next((t for t in dataset if t.ticket_id == "TKT-NONDEFAULT-003"), None)
+        self.assertIsNotNone(ticket)
+
+        env = _make_env()
+        with patch.object(env, "_dataset", [ticket]):
+            with patch.object(env, "_tickets_by_id", {ticket.ticket_id: ticket}):
+                obs = env.reset(seed=0, task_id=3, queue_size=1)
+
+        self.assertNotIn("ambiguity_note", obs.current_ticket)
+        obs = env.step(
+            HelpdeskTicketAction(
+                action_type="investigate",
+                tool_name="lookup_internal_routing_note",
+            )
+        )
+        self.assertEqual(obs.last_tool_result["routing_note"], ticket.ambiguity_note)
+        self.assertEqual(obs.current_ticket["ambiguity_note"], ticket.ambiguity_note)
+        self.assertGreater(obs.reward or 0.0, 0.0)
+
+    def test_submit_without_required_investigation_gets_shaping_penalty(self) -> None:
+        from unittest.mock import patch
+
+        dataset = load_dataset()
+        ticket = next((t for t in dataset if t.ticket_id == "TKT-NONDEFAULT-003"), None)
+        self.assertIsNotNone(ticket)
+
+        env = _make_env()
+        with patch.object(env, "_dataset", [ticket]):
+            with patch.object(env, "_tickets_by_id", {ticket.ticket_id: ticket}):
+                obs = env.reset(seed=0, task_id=3, queue_size=1)
+
+        final_obs = env.step(
+            HelpdeskTicketAction(
+                issue_type=ticket.issue_type,
+                priority=ticket.priority,
+                assignment_group=ticket.assignment_group,
+                resolution_action=ticket.resolution_action,
+            )
+        )
+
+        self.assertTrue(final_obs.done)
+        self.assertIsNotNone(final_obs.rubric_reward)
+        self.assertLess(final_obs.reward, final_obs.rubric_reward)
+        self.assertGreater(
+            final_obs.last_reward_components.get("context_gap_penalty", 0.0),
+            0.0,
+        )
+
 
 class TestQueueEconomics(unittest.TestCase):
     """Free investigations are allowed, but excessive investigation gets a queue-level penalty."""

tests/test_inference_unit.py

@@ -140,6 +140,16 @@ class InferenceUnitTests(unittest.TestCase):
         self.assertIsNone(inference.HF_TOKEN)
         self.assertFalse(inference.llm_mode_enabled())
 
+    def test_seed_env_override_is_respected(self) -> None:
+        inference = _load_inference_module({"SEED": "7"})
+
+        self.assertEqual(inference.SEED, 7)
+
+    def test_invalid_seed_env_falls_back_to_default(self) -> None:
+        inference = _load_inference_module({"SEED": "not-an-int"})
+
+        self.assertEqual(inference.SEED, 42)
+
     def test_run_uses_only_structured_start_step_end_logs(self) -> None:
         inference = _load_inference_module()
 
@@ -179,6 +189,311 @@ class InferenceUnitTests(unittest.TestCase):
             [1, 2, 3],
         )
 
+    def test_build_llm_user_message_includes_recent_history_feedback(self) -> None:
+        inference = _load_inference_module()
+
+        ticket = {
+            "ticket_id": "ticket-xyz",
+            "title": "Contractor onboarding blocked by access issue",
+            "requester": "pm@contractorco.com",
+            "description": "Access permissions are blocking contractor setup.",
+            "context_status": {
+                "investigation_required": True,
+                "revealed_tools": [],
+                "remaining_tools": ["lookup_internal_routing_note"],
+                "hints": ["An internal routing note may disambiguate the correct workflow."],
+            },
+            "last_tool_result": {"tool_name": "lookup_requester_history", "found": False},
+            "feedback_summary": "Ticket score=0.40; field_scores[issue_type=0.40]; reward=0.40",
+            "last_reward_components": {"ticket_score": 0.4, "final_reward": 0.4},
+            "investigation_budget_remaining": 2,
+            "average_score_so_far": 0.7,
+            "progress_fraction": 0.5,
+            "recent_history": [
+                {
+                    "ticket_id": "ticket-prev",
+                    "predicted": {"issue_type": "identity_access"},
+                    "score": 0.4,
+                    "breakdown": {"issue_type": 0.4},
+                    "penalty_reason": "extra_fields: ['assignment_group']",
+                    "feedback_summary": "Penalty applied: extra_fields: ['assignment_group']; reward=0.00",
+                    "reward_components": {"reward_kind": "step_penalty", "final_reward": 0.0},
+                }
+            ],
+            "queue_position": 2,
+            "tickets_remaining": 4,
+        }
+
+        message = inference.build_llm_user_message(
+            ticket,
+            ["issue_type"],
+            "Read the ticket and select the single best IT issue type.",
+        )
+
+        self.assertIn("Recent evaluation feedback", message)
+        self.assertIn("score=0.4", message)
+        self.assertIn("penalty_reason=extra_fields", message)
+        self.assertIn("Latest environment feedback", message)
+        self.assertIn("Context status", message)
+        self.assertIn("Latest reward components", message)
+        self.assertIn("Average score so far: 0.7", message)
+        self.assertIn("Episode progress: 0.5", message)
+        self.assertIn("Investigation budget remaining: 2", message)
+        self.assertIn("Investigation result", message)
+        self.assertIn("queue_position=2", message)
+
+    def test_build_action_backfills_missing_fields_from_heuristic(self) -> None:
+        inference = _load_inference_module()
+        inference.llm_client = object()
+
+        ticket = {
+            "ticket_id": "ticket-018",
+            "title": "Question about enterprise tier pricing",
+            "requester": "finance@urbanstack.io",
+            "description": (
+                "We're comparing your enterprise plan against two competitors. "
+                "Can you send over a detailed pricing breakdown?"
+            ),
+        }
+
+        with mock.patch.object(
+            inference,
+            "call_llm",
+            return_value={"issue_type": "service_request"},
+        ):
+            action, action_source, fallback_reason = inference.build_action(
+                ticket,
+                ["issue_type", "priority", "assignment_group", "resolution_action"],
+                "Perform full helpdesk routing.",
+            )
+
+        self.assertEqual(action.issue_type, "service_request")
+        self.assertEqual(action.priority, "medium")
+        self.assertEqual(action.assignment_group, "procurement")
+        self.assertEqual(action.resolution_action, "assign")
+        self.assertEqual(action_source, "llm_backfilled")
+        self.assertIn("heuristic_backfill", fallback_reason or "")
+
+    def test_build_action_ignores_invalid_llm_fields_and_keeps_valid_ones(self) -> None:
+        inference = _load_inference_module()
+        inference.llm_client = object()
+
+        ticket = {
+            "ticket_id": "ticket-018",
+            "title": "Question about enterprise tier pricing",
+            "requester": "finance@urbanstack.io",
+            "description": (
+                "We're comparing your enterprise plan against two competitors. "
+                "Can you send over a detailed pricing breakdown?"
+            ),
+        }
+
+        with mock.patch.object(
+            inference,
+            "call_llm",
+            return_value={
+                "issue_type": "service_request",
+                "priority": "urgent",
+            },
+        ):
+            action, action_source, fallback_reason = inference.build_action(
+                ticket,
+                ["issue_type", "priority"],
+                "Read the ticket, select the best IT issue type, and estimate the priority.",
+            )
+
+        self.assertEqual(action.issue_type, "service_request")
+        self.assertEqual(action.priority, "medium")
+        self.assertEqual(action_source, "llm_backfilled")
+        self.assertIn("invalid_llm_fields=['priority']", fallback_reason or "")
+
+    def test_build_action_backfills_dependent_fields_from_llm_issue_type(self) -> None:
+        inference = _load_inference_module()
+        inference.llm_client = object()
+
+        ticket = {
+            "ticket_id": "ticket-002",
+            "title": "Can not sign in after 2FA reset",
+            "requester": "ops@laneeight.io",
+            "description": (
+                "I was forced to reset 2FA and now the account stays locked even "
+                "with the backup code."
+            ),
+        }
+
+        with mock.patch.object(
+            inference,
+            "call_llm",
+            return_value={"issue_type": "identity_access"},
+        ):
+            action, action_source, fallback_reason = inference.build_action(
+                ticket,
+                ["issue_type", "assignment_group", "resolution_action"],
+                "Perform full helpdesk routing.",
+            )
+
+        self.assertEqual(action.issue_type, "identity_access")
+        self.assertEqual(action.assignment_group, "service_desk")
+        self.assertEqual(action.resolution_action, "fulfill")
+        self.assertEqual(action_source, "llm_backfilled")
+        self.assertIn("heuristic_backfill", fallback_reason or "")
+
+    def test_build_action_normalizes_pricing_request_issue_type(self) -> None:
+        inference = _load_inference_module()
+        inference.llm_client = object()
+
+        ticket = {
+            "ticket_id": "ticket-018",
+            "title": "Question about enterprise tier pricing",
+            "requester": "finance@urbanstack.io",
+            "description": (
+                "We're comparing your enterprise plan against two competitors. "
+                "Can you send over a detailed pricing breakdown?"
+            ),
+        }
+
+        with mock.patch.object(
+            inference,
+            "call_llm",
+            return_value={
+                "issue_type": "billing_license",
+                "priority": "medium",
+            },
+        ):
+            action, action_source, fallback_reason = inference.build_action(
+                ticket,
+                ["issue_type", "priority", "assignment_group", "resolution_action"],
+                "Perform full helpdesk routing.",
+            )
+
+        self.assertEqual(action.issue_type, "service_request")
+        self.assertEqual(action.assignment_group, "procurement")
+        self.assertEqual(action.resolution_action, "assign")
+        self.assertEqual(action.priority, "medium")
+        self.assertEqual(action_source, "llm_backfilled")
+        self.assertIn("domain_overrides", fallback_reason or "")
+
+    def test_build_action_normalizes_onboarding_access_blocker(self) -> None:
+        inference = _load_inference_module()
+        inference.llm_client = object()
+
+        ticket = {
+            "ticket_id": "TKT-NONDEFAULT-003",
+            "title": "Contractor onboarding blocked by access issue",
+            "requester": "pm@contractorco.com",
+            "description": (
+                "A new contractor cannot complete onboarding because their account "
+                "access is blocked by a permissions error. The onboarding team "
+                "cannot resolve access issues; routing to service desk."
+            ),
+            "ambiguity_note": "Contractor onboarding blocked by access issue, routed to service desk",
+        }
+
+        with mock.patch.object(
+            inference,
+            "call_llm",
+            return_value={
+                "issue_type": "identity_access",
+                "priority": "high",
+            },
+        ):
+            action, action_source, fallback_reason = inference.build_action(
+                ticket,
+                ["issue_type", "priority", "assignment_group", "resolution_action"],
+                "Perform full helpdesk routing.",
+            )
+
+        self.assertEqual(action.issue_type, "onboarding")
+        self.assertEqual(action.priority, "medium")
+        self.assertEqual(action.assignment_group, "service_desk")
+        self.assertEqual(action.resolution_action, "fulfill")
+        self.assertEqual(action_source, "llm_backfilled")
+        self.assertIn("domain_overrides", fallback_reason or "")
+
+    def test_build_action_deescalates_nonurgent_onboarding_priority(self) -> None:
+        inference = _load_inference_module()
+        inference.llm_client = object()
+
+        ticket = {
+            "ticket_id": "ticket-008",
+            "title": "Kickoff onboarding session for newly activated account",
+            "requester": "admin@brightpath.io",
+            "description": (
+                "We activated our account this week and need an onboarding call plus "
+                "admin setup guidance for six internal users."
+            ),
+        }
+
+        with mock.patch.object(
+            inference,
+            "call_llm",
+            return_value={
+                "issue_type": "onboarding",
+                "priority": "high",
+            },
+        ):
+            action, action_source, fallback_reason = inference.build_action(
+                ticket,
+                ["issue_type", "priority"],
+                "Read the ticket, select the best IT issue type, and estimate the priority.",
+            )
+
+        self.assertEqual(action.issue_type, "onboarding")
+        self.assertEqual(action.priority, "medium")
+        self.assertEqual(action_source, "llm_backfilled")
+        self.assertIn("domain_overrides", fallback_reason or "")
+
+    def test_merge_ticket_context_carries_feedback_summary_from_observation(self) -> None:
+        inference = _load_inference_module()
+
+        observation = SimpleNamespace(
+            last_tool_result={"tool_name": "lookup_requester_history", "found": True},
+            history=[{"ticket_id": "ticket-prev", "score": 0.4}],
+            queue_position=2,
+            tickets_remaining=4,
+            investigation_budget_remaining=1,
+            average_score_so_far=0.55,
+            progress_fraction=0.4,
+            last_reward_components={"ticket_score": 0.4, "final_reward": 0.4},
+            metadata={"last_feedback_summary": "Ticket score=0.40; reward=0.40"},
+        )
+
+        merged = inference.merge_ticket_context(
+            {
+                "ticket_id": "ticket-xyz",
+                "title": "Contractor onboarding blocked by access issue",
+            },
+            observation,
+        )
+
+        self.assertEqual(merged["feedback_summary"], "Ticket score=0.40; reward=0.40")
+        self.assertEqual(merged["investigation_budget_remaining"], 1)
+        self.assertEqual(merged["average_score_so_far"], 0.55)
+        self.assertEqual(merged["progress_fraction"], 0.4)
+        self.assertEqual(merged["last_reward_components"]["final_reward"], 0.4)
+        self.assertEqual(merged["queue_position"], 2)
+        self.assertEqual(merged["tickets_remaining"], 4)
+        self.assertEqual(merged["last_tool_result"]["tool_name"], "lookup_requester_history")
+
+    def test_should_investigate_uses_remaining_tools_from_context_status(self) -> None:
+        inference = _load_inference_module()
+
+        investigate, tool_name = inference.should_investigate(
+            {
+                "ticket_id": "ticket-021",
+                "context_status": {
+                    "remaining_tools": [
+                        "lookup_related_ticket",
+                        "lookup_requester_history",
+                    ]
+                },
+            },
+            [],
+        )
+
+        self.assertTrue(investigate)
+        self.assertEqual(tool_name, "lookup_related_ticket")
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_policy_learning.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import os
+import sys
+import types as _types
+import unittest
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+import openenv_test_stubs  # noqa: F401
+
+if "openenv.core.env_server.interfaces" not in sys.modules:
+    _interfaces_mod = _types.ModuleType("openenv.core.env_server.interfaces")
+
+    class _Environment:
+        def __init__(self) -> None:
+            pass
+
+        def __init_subclass__(cls, **kwargs: object) -> None:
+            super().__init_subclass__(**kwargs)
+
+        @classmethod
+        def __class_getitem__(cls, item: object) -> type:
+            return cls
+
+    _interfaces_mod.Environment = _Environment  # type: ignore[attr-defined]
+    sys.modules["openenv.core.env_server.interfaces"] = _interfaces_mod
+
+
+from models import HelpdeskTicketAction, HelpdeskTicketObservation
+from policy_learning import (
+    POLICY_LIBRARY,
+    choose_policy_action,
+    compare_policies,
+    parse_int_spec,
+    rollout_episode,
+    search_policies,
+)
+from server.environment import HelpdeskTicketRoutingEnvironment
+from server.tasks import get_task_definition
+
+
+class SingleTicketEnvironment(HelpdeskTicketRoutingEnvironment):
+    def __init__(self, ticket_id: str) -> None:
+        super().__init__()
+        self._forced_ticket_id = ticket_id
+
+    def reset(self, seed=None, episode_id=None, **kwargs):
+        observation = super().reset(seed=seed, episode_id=episode_id, **kwargs)
+        ticket = self._tickets_by_id[self._forced_ticket_id]
+        self._queue = [ticket]
+        self._state.current_task_id = int(kwargs.get("task_id", 3))
+        self._state.queue_ticket_ids = [ticket.ticket_id]
+        self._state.current_ticket_index = 0
+        self._state.per_ticket_scores = []
+        self._state.total_reward = 0.0
+        self._state.last_step_reward = None
+        self._state.reward = None
+        self._state.done = False
+        self._state.average_score_so_far = 0.0
+        self._state.investigation_steps = 0
+        self._state.investigation_budget_remaining = len(self._queue)
+        self._state.investigation_penalty_applied = 0.0
+        self._state.last_tool_result = None
+        self._state.last_reward_components = {}
+        self._state.ticket_tool_usage = {}
+        self._state.history_entries = []
+        return self._build_observation(get_task_definition(self._state.current_task_id))
+
+
+def _context_sensitive_submit_builder(
+    ticket: dict[str, object], allowed_fields: list[str]
+) -> HelpdeskTicketAction:
+    if ticket.get("ambiguity_note"):
+        values = {
+            "issue_type": "onboarding",
+            "priority": "medium",
+            "assignment_group": "service_desk",
+            "resolution_action": "fulfill",
+        }
+    else:
+        values = {
+            "issue_type": "identity_access",
+            "priority": "high",
+            "assignment_group": "service_desk",
+            "resolution_action": "fulfill",
+        }
+    return HelpdeskTicketAction(
+        **{field: value for field, value in values.items() if field in allowed_fields}
+    )
+
+
+class PolicyLearningTests(unittest.TestCase):
+    def test_parse_int_spec_expands_ranges(self) -> None:
+        self.assertEqual(parse_int_spec("42-44,44,46", field_name="seeds"), [42, 43, 44, 46])
+
+    def test_choose_policy_action_prefers_hidden_context_tools(self) -> None:
+        policy = POLICY_LIBRARY["investigate_when_context_hidden"]
+        observation = HelpdeskTicketObservation(
+            current_ticket={
+                "ticket_id": "ticket-021",
+                "context_status": {
+                    "remaining_tools": ["lookup_related_ticket", "lookup_requester_history"],
+                    "revealed_tools": [],
+                }
+            },
+            allowed_fields=["issue_type"],
+        )
+
+        action, source = choose_policy_action(policy, observation, {}, _context_sensitive_submit_builder)
+
+        self.assertEqual(action.action_type, "investigate")
+        self.assertEqual(action.tool_name, "lookup_related_ticket")
+        self.assertEqual(source, "investigate_hidden_context")
+
+    def test_choose_policy_action_submits_when_investigation_disabled(self) -> None:
+        policy = POLICY_LIBRARY["no_investigation"]
+        observation = HelpdeskTicketObservation(
+            current_ticket={
+                "ticket_id": "ticket-021",
+                "context_status": {"remaining_tools": ["lookup_related_ticket"]},
+            },
+            allowed_fields=["issue_type", "priority"],
+        )
+
+        action, source = choose_policy_action(policy, observation, {}, _context_sensitive_submit_builder)
+
+        self.assertEqual(action.action_type, "submit")
+        self.assertEqual(action.issue_type, "identity_access")
+        self.assertEqual(source, "submit")
+
+    def test_rollout_episode_rewards_context_aware_policy(self) -> None:
+        no_investigation = POLICY_LIBRARY["no_investigation"]
+        context_aware = POLICY_LIBRARY["investigate_when_context_hidden"]
+
+        no_summary, _ = rollout_episode(
+            env=SingleTicketEnvironment("TKT-NONDEFAULT-003"),
+            policy=no_investigation,
+            seed=42,
+            task_id=3,
+            submit_builder=_context_sensitive_submit_builder,
+        )
+        context_summary, _ = rollout_episode(
+            env=SingleTicketEnvironment("TKT-NONDEFAULT-003"),
+            policy=context_aware,
+            seed=42,
+            task_id=3,
+            submit_builder=_context_sensitive_submit_builder,
+        )
+
+        self.assertLess(no_summary["terminal_reward"], context_summary["terminal_reward"])
+        self.assertLess(no_summary["normalized_return"], context_summary["normalized_return"])
+        self.assertEqual(context_summary["investigation_steps"], 1)
+
+    def test_search_policies_selects_better_policy(self) -> None:
+        report = search_policies(
+            [
+                POLICY_LIBRARY["no_investigation"],
+                POLICY_LIBRARY["investigate_when_context_hidden"],
+            ],
+            train_seeds=[41, 42],
+            eval_seeds=[43],
+            task_ids=[3],
+            output_dir=os.path.join(os.getcwd(), "analysis", "policy_learning_test"),
+            env_factory=lambda: SingleTicketEnvironment("TKT-NONDEFAULT-003"),
+            submit_builder=_context_sensitive_submit_builder,
+        )
+
+        self.assertEqual(report["selected_policy"], "investigate_when_context_hidden")
+        self.assertGreater(
+            report["eval_improvement_vs_baseline"]["avg_normalized_return"],
+            0.0,
+        )
+
+    def test_compare_policies_reports_improvement(self) -> None:
+        report = compare_policies(
+            [
+                POLICY_LIBRARY["no_investigation"],
+                POLICY_LIBRARY["investigate_when_context_hidden"],
+            ],
+            seeds=[42],
+            task_ids=[3],
+            output_dir=os.path.join(os.getcwd(), "analysis", "policy_learning_compare_test"),
+            env_factory=lambda: SingleTicketEnvironment("TKT-NONDEFAULT-003"),
+            submit_builder=_context_sensitive_submit_builder,
+        )
+
+        self.assertEqual(report["best_policy"], "investigate_when_context_hidden")
+        self.assertGreater(report["improvement_vs_baseline"]["avg_terminal_reward"], 0.0)
+
+
+if __name__ == "__main__":
+    unittest.main()