""" LedgerShield OpenEnv Environment. Implements a POMDP-style environment for evaluating AI agents on enterprise accounts-payable (AP) payment integrity tasks. The agent must investigate invoices, gather evidence, take interventions, and submit a final payment decision. Environment Loop: 1. ``reset()`` loads a case and returns the initial observation. 2. ``step(action)`` processes one action (tool call, intervention, or final decision) and returns the next observation. 3. The episode ends when the agent submits a decision, exhausts its budget, or exceeds the maximum step count. Reward Design: - Potential-Based Reward Shaping (PBRS) with configurable scale. - Information-gain bonus for discovering novel risk signals. - Milestone rewards for completing key investigation steps. - Terminal reward from the grading rubric. Gymnasium Compatibility: - ``truncated`` vs ``terminated`` distinction (3.2). - ``render()`` method for text-based episode summaries (3.3). - ``action_space()`` / ``observation_space()`` class methods (3.4). """ from __future__ import annotations from copy import deepcopy import math from dataclasses import asdict import os import random import re import uuid from typing import Any from models import LedgerShieldObservation, LedgerShieldReward, LedgerShieldState from openenv_compat import Environment from .compliance_engine import evaluate_compliance from .currency_engine import validate_iban, validate_swift from .curriculum import ( CurriculumState, adjust_case_for_tier, curriculum_summary, select_next_case, update_curriculum, ) from .data_loader import load_all from .dual_agent_mode import ( StackelbergAuditStrategy, WatchdogState, build_watchdog_observation, compute_stackelberg_equilibrium, score_dual_agent_episode, update_watchdog_state, watchdog_evaluate_decision, ) from .reward_machine import ( RewardMachineState, initialize_reward_machine, reward_machine_payload, transition_reward_machine, ) from .categorical_composition import task_family_component from .rl_export import export_state_vector from .benchmark_contract import ( BLIND_CONTROL_TRACK, CASE_TRACK, case_matches_track, case_track_metadata, normalize_track, track_description, track_label, ) from .control_statechart import control_boundary_snapshot, evaluate_control_boundary from .grading import score_submission from .decision_certificate import build_decision_certificate, verify_decision_certificate from .decision_falsifier import falsify_decision from .institutional_game import ( InstitutionalMemory, attach_institutional_context, evaluate_authority_gate, institutional_context_for_case, public_institutional_memory, record_trust_graph, record_institutional_outcome, ) from .outcome_simulator import simulate_outcome from .proper_scoring import resolve_predicted_probabilities from .risk_rules import assess_submission_risk from .schema import ALLOWED_ACTIONS, ALLOWED_DECISIONS, INTERVENTION_ACTIONS from .sprt_engine import ( DEFAULT_HYPOTHESES, SPRTState, infer_tool_observation, initialize_sprt, optimal_stopping_check, sprt_state_payload, update_sprt, ) from .tools import ( compare_bank_account_tool, get_doc_crop_tool, inspect_email_thread_tool, lookup_po_tool, lookup_policy_tool, lookup_receipt_tool, lookup_vendor_history_tool, lookup_vendor_tool, ocr_tool, search_ledger_tool, zoom_tool, ) from .transition_engine import handle_intervention, normalized_result_with_signals from .trust_graph import build_trust_graph from .voi_engine import optimal_tool_selection, value_of_information from .world_state import ( advance_pending_events, build_hidden_world, decision_readiness, inject_pressure_event, investigation_status, pending_events_public, pressure_resistance_score, public_state_snapshot, public_revealed_artifacts, risk_snapshot, state_potential, system_state_snapshot, ) # ── Tool cost table ────────────────────────────────────────────────────────── TOOL_COSTS = { "zoom": 0.20, "get_doc_crop": 0.20, "ocr_fast": 0.45, "ocr_accurate": 1.10, "lookup_vendor": 0.20, "lookup_vendor_history": 0.25, "lookup_policy": 0.15, "lookup_po": 0.20, "lookup_receipt": 0.20, "search_ledger": 0.35, "inspect_email_thread": 0.25, "compare_bank_account": 0.15, "request_callback_verification": 0.40, "freeze_vendor_profile": 0.20, "request_bank_change_approval_chain": 0.30, "request_po_reconciliation": 0.30, "request_additional_receipt_evidence": 0.25, "route_to_procurement": 0.15, "route_to_security": 0.20, "flag_duplicate_cluster_review": 0.25, "create_human_handoff": 0.20, "submit_decision": 0.0, } # ── Reward shaping constants (Phase 3.1) ───────────────────────────────────── SHAPING_GAMMA = 0.98 SHAPING_SCALE = 0.35 # Upgraded from 0.18 → 0.35 # ── Information-gain bonus (Phase 5.3) ─────────────────────────────────────── INFO_GAIN_BONUS = 0.08 # ── Milestone reward definitions ───────────────────────────────────────────── MILESTONE_REWARDS: dict[str, float] = { "first_risk_signal": 0.05, "callback_requested": 0.04, "all_required_actions": 0.06, "artifact_revealed": 0.03, } # ── Degenerate evidence cap (Phase 4.5) ────────────────────────────────────── DEGENERATE_EVIDENCE_CAP = 0.25 _CUSTOM_CASE_ID_RE = re.compile(r"^CUSTOM-[A-Z0-9]{2,16}$") # ── Formalized score constants ─────────────────────────────────────────────── INTERVENTION_BASE_SCORE = 0.15 # Tightened from 0.35 (Phase 2.3) class LedgerShieldEnvironment(Environment): """POMDP environment for enterprise payment integrity evaluation. This environment simulates a realistic accounts-payable investigation workflow where an AI agent must analyze invoices, verify vendor identities, check policies, and make payment decisions. The agent operates under partial observability: it cannot see hidden risk signals directly but must discover them through tool usage and interventions. Attributes: db: Pre-loaded database of cases, vendors, policies, etc. rng: Seeded random number generator. current_case: The currently loaded case dictionary. """ def __init__(self, db: dict[str, Any] | None = None) -> None: """Initialize the LedgerShield environment. Args: db: Optional pre-loaded database dict. If None, loads from fixture files via ``load_all()``. """ super().__init__() self.db = db if db is not None else load_all() self.rng = random.Random(42) self.current_case: dict[str, Any] | None = None self._state = LedgerShieldState() self._last_reward = 0.0 self._last_done = False self._last_truncated = False self._last_terminated = False self._last_info: dict[str, Any] = {} self._hidden_world: dict[str, Any] = {} self._milestones_awarded: set[str] = set() self._render_mode: str | None = None self._curriculum_state = CurriculumState() self._watchdog_state = WatchdogState() self._sprt_runtime_state: SPRTState = initialize_sprt() self._reward_machine_runtime_state: RewardMachineState = initialize_reward_machine("task_a") self._institutional_memory = InstitutionalMemory.from_cases(self.db.get("cases", [])) self._track_mode = os.getenv("LEDGERSHIELD_TRACK_MODE", "blind").strip().lower() or "blind" self._benchmark_track = CASE_TRACK # ── Gymnasium-compatible space definitions (Phase 3.4) ─────────────── @classmethod def action_space(cls) -> dict[str, Any]: """Return a formal description of the action space. The action space is a dictionary with: - ``type``: ``"Dict"`` (composite action). - ``action_type``: ``"Discrete"`` over allowed action strings. - ``payload``: ``"Dict"`` with tool-specific parameters. Returns: Dictionary describing the action space structure. """ return { "type": "Dict", "spaces": { "action_type": { "type": "Discrete", "n": len(ALLOWED_ACTIONS), "values": list(ALLOWED_ACTIONS), }, "payload": { "type": "Dict", "description": "Tool-specific parameters (varies by action_type)", "examples": { "zoom": {"doc_id": "str", "page": "int", "region": "[x1,y1,x2,y2]"}, "ocr": {"doc_id": "str", "mode": "'fast'|'accurate'"}, "lookup_vendor": {"vendor_key": "str"}, "submit_decision": { "decision": "PAY|HOLD|NEEDS_REVIEW|ESCALATE_FRAUD", "confidence": "float(0-1)", "reason_codes": "list[str]", "predicted_probabilities": "dict[hypothesis,float] (optional)", "decision_certificate": "Decision Certificate Graph (optional)", }, }, }, }, } @classmethod def observation_space(cls) -> dict[str, Any]: """Return a formal description of the observation space. Returns: Dictionary describing the observation space structure. """ return { "type": "Dict", "spaces": { "case_id": {"type": "Text"}, "task_type": {"type": "Discrete", "values": ["task_a", "task_b", "task_c", "task_d", "task_e"]}, "instruction": {"type": "Text"}, "visible_documents": {"type": "Sequence", "element": "DocumentCatalogEntry"}, "revealed_artifacts": {"type": "Sequence", "element": "ArtifactEntry"}, "pending_events": {"type": "Sequence", "element": "PendingEvent"}, "budget_remaining": {"type": "Box", "low": 0.0, "high": 30.0}, "budget_total": {"type": "Box", "low": 0.0, "high": 30.0}, "step_count": {"type": "Discrete", "low": 0, "high": 50}, "max_steps": {"type": "Discrete", "low": 1, "high": 50}, "case_clock": {"type": "Discrete", "low": 0, "high": 50}, "risk_snapshot": {"type": "Dict"}, "investigation_status": {"type": "Dict"}, "last_tool_result": {"type": "Dict"}, "messages": {"type": "Sequence", "element": "Text"}, "allowed_actions": {"type": "Sequence", "element": "Text"}, "available_interventions": {"type": "Sequence", "element": "Text"}, "case_metadata": {"type": "Dict"}, "portfolio_context": {"type": "Dict"}, "sprt_state": {"type": "Dict"}, "tool_rankings": {"type": "Dict"}, "reward_machine": {"type": "Dict"}, "institutional_memory": {"type": "Dict"}, }, } # ── Properties ─────────────────────────────────────────────────────── @property def state(self) -> LedgerShieldState: """Return the current internal state.""" return self._state def public_state(self) -> dict[str, Any]: """Return the public (non-hidden) state snapshot.""" state = public_state_snapshot(self._state, self._hidden_world) state["institutional_memory"] = public_institutional_memory(self._institutional_memory) state["control_boundary"] = control_boundary_snapshot(self._state, self._hidden_world) return state def institutional_memory(self) -> dict[str, Any]: """Return the persistent institutional memory/loss ledger.""" return public_institutional_memory(self._institutional_memory) def reset_institutional_memory(self) -> dict[str, Any]: """Reset persistent portfolio memory without changing fixture data.""" self._institutional_memory = InstitutionalMemory.from_cases(self.db.get("cases", [])) return self.institutional_memory() # ── Internal helpers ───────────────────────────────────────────────── def _normalize_custom_case_payload( self, raw: Any, ) -> dict[str, Any]: """Validate API ``custom_case`` and return normalized fields for cloning. Clones an existing benchmark case (template) and assigns a new ``case_id`` plus instruction text. Document graphs and gold stay identical to the template. """ if not isinstance(raw, dict): raise ValueError("custom_case must be a JSON object") template_id = raw.get("template_case_id") if template_id is None: template_id = raw.get("templateCaseId") if not isinstance(template_id, str) or not template_id.strip(): raise ValueError("custom_case.template_case_id is required") template_id = template_id.strip() cases_by_id = self.db.get("cases_by_id") or {} if template_id not in cases_by_id: raise ValueError(f"unknown custom_case.template_case_id: {template_id}") case_id = raw.get("case_id") if case_id is None: case_id = raw.get("caseId") if not isinstance(case_id, str) or not case_id.strip(): raise ValueError("custom_case.case_id is required") case_id_norm = case_id.strip().upper() if not _CUSTOM_CASE_ID_RE.match(case_id_norm): raise ValueError( "custom_case.case_id must match CUSTOM- plus 2-16 uppercase letters/digits " "(example: CUSTOM-DEMO01)" ) instruction = raw.get("instruction") if not isinstance(instruction, str) or not instruction.strip(): raise ValueError("custom_case.instruction is required (non-empty string)") instruction = instruction.strip() if "\n" in instruction or "\r" in instruction: raise ValueError("custom_case.instruction must be a single line (no newlines)") if len(instruction) > 800: raise ValueError("custom_case.instruction must be at most 800 characters") out: dict[str, Any] = { "template_case_id": template_id, "case_id": case_id_norm, "instruction": instruction, } if raw.get("max_steps") is not None or raw.get("maxSteps") is not None: ms = raw.get("max_steps", raw.get("maxSteps")) if isinstance(ms, bool): raise ValueError("custom_case.max_steps must be an integer") if isinstance(ms, int): ms_val = ms elif isinstance(ms, float) and ms.is_integer(): ms_val = int(ms) else: raise ValueError("custom_case.max_steps must be an integer") if ms_val < 4 or ms_val > 50: raise ValueError("custom_case.max_steps must be between 4 and 50") out["max_steps"] = ms_val if raw.get("budget_total") is not None or raw.get("budgetTotal") is not None: bt = raw.get("budget_total", raw.get("budgetTotal")) if isinstance(bt, (int, float)): btf = float(bt) else: raise ValueError("custom_case.budget_total must be a number") if btf < 1.0 or btf > 50.0: raise ValueError("custom_case.budget_total must be between 1 and 50") out["budget_total"] = round(btf, 4) return out def _case_from_custom_spec(self, spec: dict[str, Any]) -> dict[str, Any]: template = self.db["cases_by_id"][spec["template_case_id"]] case = deepcopy(template) case["case_id"] = spec["case_id"] case["instruction"] = spec["instruction"] if "max_steps" in spec: case["max_steps"] = spec["max_steps"] if "budget_total" in spec: case["budget_total"] = spec["budget_total"] return case def _select_case( self, seed: int | None = None, case_id: str | None = None, track: str | None = None, ) -> dict[str, Any]: """Select a case by ID or random sampling. Args: seed: Random seed for case selection. case_id: Specific case ID to load. Returns: Case dictionary. Raises: ValueError: If case_id is provided but not found. """ if case_id: case = self.db["cases_by_id"].get(case_id) if case is None: raise ValueError(f"unknown case_id: {case_id}") return case selection_seed = seed if seed is not None else self.rng.randint(0, 2**31 - 1) requested_track = normalize_track(track) candidate_cases = [ case for case in self.db["cases"] if case_matches_track(case, requested_track) ] or list(self.db["cases"]) selected = select_next_case(self._curriculum_state, candidate_cases, seed=selection_seed) return adjust_case_for_tier(selected, self._curriculum_state.tier) def _currency_validation_snapshot(self, submitted: dict[str, Any]) -> dict[str, Any]: assert self.current_case is not None task_type = str(self.current_case.get("task_type", "")) if task_type != "task_a": return {"applicable": False, "score": 1.0} gold_fields = (self.current_case.get("gold", {}) or {}).get("fields", {}) or {} extracted_fields = submitted.get("extracted_fields", {}) or {} expected_bank = str(gold_fields.get("bank_account", "") or "").strip() submitted_bank = str(extracted_fields.get("bank_account", "") or "").strip() expected_currency = str(gold_fields.get("currency", "") or "").strip().upper() submitted_currency = str(extracted_fields.get("currency", "") or "").strip().upper() checks: list[float] = [] snapshot: dict[str, Any] = {"applicable": True, "format": "unknown"} if expected_currency: checks.append(float(submitted_currency == expected_currency)) snapshot["expected_currency"] = expected_currency snapshot["submitted_currency"] = submitted_currency compact_expected = "".join(expected_bank.split()).upper() if expected_bank: checks.append(float(" ".join(submitted_bank.lower().split()) == " ".join(expected_bank.lower().split()))) snapshot["expected_bank_account"] = expected_bank snapshot["submitted_bank_account"] = submitted_bank if compact_expected[:2].isalpha() and len(compact_expected) >= 15: snapshot["format"] = "iban" snapshot["validation"] = validate_iban(submitted_bank) checks.append(float(snapshot["validation"].get("valid", False))) elif len(compact_expected) in {8, 11} and compact_expected[:4].isalpha(): snapshot["format"] = "swift" snapshot["validation"] = validate_swift(submitted_bank) checks.append(float(snapshot["validation"].get("valid", False))) snapshot["score"] = round(sum(checks) / len(checks), 4) if checks else 1.0 if not checks: snapshot["applicable"] = False return snapshot def _initial_visible_doc_ids(self) -> list[str]: """Return initial visible document IDs for the current case.""" assert self.current_case is not None doc_ids = self.current_case.get("initial_visible_doc_ids") or [ doc.get("doc_id") for doc in self.current_case.get("documents", []) if doc.get("doc_id") ] return [str(doc_id) for doc_id in doc_ids] def _all_documents(self) -> list[dict[str, Any]]: """Return all documents (static + dynamic) for the current case.""" assert self.current_case is not None docs = list(self.current_case.get("documents", [])) dynamic_docs = self._hidden_world.get("dynamic_documents", {}) or {} docs.extend(dynamic_docs.values()) return docs def _visible_document_catalog(self) -> list[dict[str, Any]]: """Build the visible document catalog for the current observation.""" assert self.current_case is not None docs: list[dict[str, Any]] = [] visible_set = set(self._state.visible_doc_ids) for doc in self._all_documents(): doc_id = str(doc.get("doc_id")) if doc_id not in visible_set: continue docs.append( { "doc_id": doc_id, "doc_type": doc.get("doc_type", "unknown"), "thumbnail": doc.get("thumbnail", f"thumbnail::{doc_id}"), "page_count": doc.get("page_count", 1), "language": doc.get("language", "en"), "available_views": [ "thumbnail", "zoom", "get_doc_crop", "ocr_fast", "ocr_accurate", ], } ) return docs def _observation( self, tool_result: dict[str, Any] | None = None, messages: list[str] | None = None, extra_metadata: dict[str, Any] | None = None, ) -> LedgerShieldObservation: """Construct an observation from the current state. Args: tool_result: Result of the last tool call (if any). messages: List of messages to include in the observation. extra_metadata: Additional key-value pairs merged into case_metadata. Used by reset() to expose the MDPComponent categorical spec. Returns: LedgerShieldObservation dataclass. """ assert self.current_case is not None base_metadata: dict[str, Any] = { "task_label": self.current_case.get("task_label", ""), "due_date_days": int(self.current_case.get("due_date_days", 14) or 14), "ashtg": "Adversarial Sequential Hypothesis Testing Game", "benchmark_identity": "Verified institutional control intelligence in enterprise AP workflows", "benchmark_track": self._benchmark_track, "benchmark_track_label": track_label(self._benchmark_track), "benchmark_track_description": track_description(self._benchmark_track), "official_tracks": list(self.current_case.get("official_tracks", [])), } if extra_metadata: base_metadata.update(extra_metadata) observation_track_mode = "blind" if self._benchmark_track == BLIND_CONTROL_TRACK else self._track_mode base_metadata["track_mode"] = observation_track_mode instrumented = observation_track_mode != "blind" return LedgerShieldObservation( case_id=self._state.case_id, task_type=self._state.task_type, instruction=self.current_case["instruction"], visible_documents=self._visible_document_catalog(), revealed_artifacts=public_revealed_artifacts(self._state, self._hidden_world), pending_events=pending_events_public(self._hidden_world), budget_remaining=round(self._state.budget_remaining, 3), budget_total=round(self._state.budget_total, 3), step_count=self._state.step_count, max_steps=self._state.max_steps, case_clock=self._state.case_clock, risk_snapshot=risk_snapshot(self._state, self._hidden_world), investigation_status=investigation_status(self._state), last_tool_result=tool_result or {}, messages=messages or [], allowed_actions=list(ALLOWED_ACTIONS), available_interventions=list(INTERVENTION_ACTIONS), case_metadata=base_metadata, portfolio_context=dict(self._hidden_world.get("campaign_context", {})), sprt_state=deepcopy(self._state.sprt_state) if instrumented else {}, tool_rankings=deepcopy(self._state.tool_rankings) if instrumented else {}, reward_machine=deepcopy(self._state.reward_machine_state) if instrumented else {}, institutional_memory=public_institutional_memory(self._institutional_memory), ) def _reward_payload( self, *, value: float, terminal: bool, components: dict[str, float] | None = None, metadata: dict[str, Any] | None = None, ) -> dict[str, Any]: """Build a structured reward payload. Args: value: Scalar reward value. terminal: Whether this is the terminal reward. components: Breakdown of reward components. metadata: Additional reward metadata. Returns: Serialized LedgerShieldReward dict. """ return LedgerShieldReward( value=round(float(value), 4), terminal=terminal, components={key: round(float(val), 4) for key, val in (components or {}).items()}, metadata=metadata or {}, ).model_dump() # ── Milestone tracking (Phase 3.1) ─────────────────────────────────── def _check_milestones(self) -> float: """Check and award milestone rewards. Returns: Total milestone reward for this step. """ bonus = 0.0 # First risk signal discovery if (self._state.observed_risk_signals and "first_risk_signal" not in self._milestones_awarded): self._milestones_awarded.add("first_risk_signal") bonus += MILESTONE_REWARDS["first_risk_signal"] # Callback requested callback_taken = any( step.get("action_type") == "request_callback_verification" for step in self._state.trajectory ) if callback_taken and "callback_requested" not in self._milestones_awarded: self._milestones_awarded.add("callback_requested") bonus += MILESTONE_REWARDS["callback_requested"] # Artifact revealed if (self._state.revealed_artifact_ids and "artifact_revealed" not in self._milestones_awarded): self._milestones_awarded.add("artifact_revealed") bonus += MILESTONE_REWARDS["artifact_revealed"] # All required actions completed required = set(self._hidden_world.get("required_actions", [])) successful = { step.get("action_type", "") for step in self._state.trajectory if step.get("success", True) } if required and required <= successful and "all_required_actions" not in self._milestones_awarded: self._milestones_awarded.add("all_required_actions") bonus += MILESTONE_REWARDS["all_required_actions"] return bonus # ── Information-theoretic exploration bonus (Phase 5.3) ─────────────── def _info_gain_bonus(self, signals_before: int, signals_after: int) -> float: """Calculate information-gain bonus for discovering new risk signals. Uses an entropy-inspired formula: bonus scales with the log-ratio of information gained, saturating at INFO_GAIN_BONUS. Args: signals_before: Number of observed risk signals before action. signals_after: Number of observed risk signals after action. Returns: Float bonus value. """ new_signals = max(0, signals_after - signals_before) if new_signals == 0: return 0.0 total_hidden = max(len(self._hidden_world.get("hidden_risk_signals", [])), 1) coverage_before = signals_before / total_hidden coverage_after = signals_after / total_hidden # Log-ratio information gain (bounded) if coverage_before >= 1.0: return 0.0 gain = math.log2(max(1.0 - coverage_before, 0.01)) - math.log2(max(1.0 - coverage_after, 0.01)) return min(INFO_GAIN_BONUS, gain * 0.04) # ── Core API ───────────────────────────────────────────────────────── def reset( self, seed: int | None = None, case_id: str | None = None, track: str | None = None, custom_case: dict[str, Any] | None = None, ) -> LedgerShieldObservation: """Reset the environment and load a new case. Args: seed: Optional seed for case selection. case_id: Optional specific case to load. custom_case: Optional dict to clone a template case under a new ``CUSTOM-…`` id and instruction (validated); when set, ``case_id`` / ``seed`` selection for loading is ignored. Returns: Initial observation for the new episode. """ if custom_case is not None: spec = self._normalize_custom_case_payload(custom_case) self.current_case = self._case_from_custom_spec(spec) else: self.current_case = self._select_case(seed=seed, case_id=case_id, track=track) self._benchmark_track = normalize_track(track or self.current_case.get("primary_track")) self._hidden_world = build_hidden_world(self.current_case) institutional_context = institutional_context_for_case( self.current_case, self.db.get("cases", []), self._institutional_memory, ) attach_institutional_context(self._hidden_world, institutional_context) self._state = LedgerShieldState( episode_id=str(uuid.uuid4()), case_id=self.current_case["case_id"], task_type=self.current_case["task_type"], budget_total=self.current_case.get("budget_total", 15.0), budget_remaining=self.current_case.get("budget_total", 15.0), max_steps=self.current_case.get("max_steps", 20), visible_doc_ids=self._initial_visible_doc_ids(), difficulty=self.current_case.get("difficulty", "medium"), hidden_risk_signals=list(self._hidden_world.get("hidden_risk_signals", [])), portfolio_metrics=dict(self._hidden_world.get("campaign_context", {})), contrastive_pair_id=str(self.current_case.get("contrastive_pair_id", "")), ) self._last_reward = 0.0 self._last_done = False self._last_truncated = False self._last_terminated = False self._last_info = {"case_id": self._state.case_id} self._milestones_awarded = set() self._sprt_runtime_state = initialize_sprt(hypotheses=DEFAULT_HYPOTHESES) self._reward_machine_runtime_state = initialize_reward_machine(self._state.task_type) self._watchdog_state = WatchdogState(strategy=self._apply_stackelberg_strategy()) self._state.calibration_running_average = 0.0 # ── Categorical MDP Composition (Pillar 9) ──────────────────────── # Load the MDPComponent for this task family. The component defines # the task's formal state/action spaces and temporal specification. # required_observations seeds the hidden world's required_actions so # milestone detection and VoI computation know what evidence the task # demands. mdp_component = task_family_component(self._state.task_type) self._mdp_component = mdp_component # Inject the component's required action-space into the hidden world # so that _check_milestones() can verify completion against the # categorical spec rather than a hard-coded list. if "required_actions" not in self._hidden_world or not self._hidden_world["required_actions"]: self._hidden_world["required_actions"] = sorted(mdp_component.action_space) self._refresh_ashtg_public_state() self._state.decision_readiness = round(decision_readiness(self._state, self._hidden_world), 4) tier_name = curriculum_summary(self._curriculum_state).get("tier_name", "unknown") mdp_spec = { "component_name": mdp_component.name, "action_space": sorted(mdp_component.action_space), "state_space": sorted(mdp_component.state_space), "required_observations": sorted(mdp_component.required_observations), "temporal_spec": mdp_component.temporal_spec, } benchmark_metadata = case_track_metadata(self.current_case) return self._observation( messages=[f"Loaded case {self._state.case_id} (curriculum: {tier_name})"], extra_metadata={ "mdp_component": mdp_spec, "benchmark_track": self._benchmark_track, "benchmark_track_label": track_label(self._benchmark_track), "benchmark_track_description": track_description(self._benchmark_track), "official_tracks": benchmark_metadata["official_tracks"], }, ) def _apply_cost(self, tool_name: str, payload: dict[str, Any]) -> float: """Calculate the budget cost for a tool invocation. Args: tool_name: Name of the tool being called. payload: Tool payload (used for OCR mode selection). Returns: Float cost value. """ if tool_name == "ocr": return TOOL_COSTS["ocr_accurate"] if payload.get("mode", "fast") == "accurate" else TOOL_COSTS["ocr_fast"] return TOOL_COSTS.get(tool_name, 0.0) def _dispatch_tool(self, tool_name: str, payload: dict[str, Any]) -> dict[str, Any]: """Dispatch a tool call to the appropriate handler. Args: tool_name: Name of the tool to invoke. payload: Tool-specific parameters. Returns: Raw tool result dictionary. """ assert self.current_case is not None overrides = self.current_case.get("context_overrides", {}) or {} vendor_index = self.db["vendors_by_key"] override_vendors = overrides.get("vendors_by_key") or overrides.get("vendors") if isinstance(override_vendors, dict): vendor_index = override_vendors elif isinstance(override_vendors, list): vendor_index = { normalize_text(vendor.get("vendor_key")): vendor for vendor in override_vendors if isinstance(vendor, dict) and vendor.get("vendor_key") } or vendor_index po_index = self.db["po_by_id"] override_pos = overrides.get("po_by_id") or overrides.get("po_records") if isinstance(override_pos, dict): po_index = override_pos elif isinstance(override_pos, list): po_index = { str(row.get("po_id")): row for row in override_pos if isinstance(row, dict) and row.get("po_id") } or po_index receipt_index = self.db["receipt_by_id"] override_receipts = overrides.get("receipt_by_id") or overrides.get("receipts") if isinstance(override_receipts, dict): receipt_index = override_receipts elif isinstance(override_receipts, list): receipt_index = { str(row.get("receipt_id")): row for row in override_receipts if isinstance(row, dict) and row.get("receipt_id") } or receipt_index email_threads = overrides.get("email_threads", self.db["email_threads"]) dispatch_map = { "zoom": lambda: zoom_tool(self.current_case, payload), "get_doc_crop": lambda: get_doc_crop_tool(self.current_case, payload), "ocr": lambda: ocr_tool(self.current_case, payload), "lookup_vendor": lambda: lookup_vendor_tool(vendor_index, payload), "lookup_vendor_history": lambda: lookup_vendor_history_tool( overrides.get("vendor_history", self.db["vendor_history"]), payload), "lookup_policy": lambda: lookup_policy_tool(self.db["policy_by_id"], self.db["policy_rules"], payload), "lookup_po": lambda: lookup_po_tool(po_index, payload), "lookup_receipt": lambda: lookup_receipt_tool(receipt_index, payload), "search_ledger": lambda: search_ledger_tool( self.current_case, overrides.get("ledger_index", self.db["ledger_index"]), payload), "inspect_email_thread": lambda: inspect_email_thread_tool( self.current_case, email_threads, payload), "compare_bank_account": lambda: compare_bank_account_tool(vendor_index, payload), } handler = dispatch_map.get(tool_name) if handler: return handler() return {"error": f"unknown action_type: {tool_name}"} def _handle_intervention( self, action_type: str, payload: dict[str, Any], ) -> tuple[dict[str, Any], list[str]]: """Handle an intervention action. Args: action_type: The intervention action type. payload: Intervention parameters. Returns: Tuple of (result_dict, messages_list). """ return handle_intervention( state=self._state, hidden_world=self._hidden_world, action_type=action_type, payload=payload, ) def _normalize_tool_result( self, tool_name: str, raw: dict[str, Any], cost: float, ) -> tuple[dict[str, Any], list[str]]: """Normalize a raw tool result into a standard format. Args: tool_name: Name of the tool. raw: Raw result from tool dispatch. cost: Budget cost incurred. Returns: Tuple of (normalized_result, messages). """ return normalized_result_with_signals( state=self._state, tool_name=tool_name, raw=raw, cost=cost, ) def _investigation_summary(self) -> dict[str, Any]: """Build a summary of the investigation for grading. Returns: Dictionary with investigation statistics. """ return { "tool_calls": len(self._state.tool_trace), "interventions_taken": len(self._state.interventions_taken), "revealed_artifact_ids": list(self._state.revealed_artifact_ids), "observed_risk_signals": list(self._state.observed_risk_signals), "sprt_recommendation": (self._state.sprt_state or {}).get("recommended_decision"), } def _voi_channel_for_action(self, action_type: str) -> str: return { "request_callback_verification": "callback_verification_result", "flag_duplicate_cluster_review": "duplicate_cluster_report", "request_bank_change_approval_chain": "bank_change_approval_chain", "request_po_reconciliation": "po_reconciliation_report", "request_additional_receipt_evidence": "receipt_reconciliation_report", }.get(action_type, action_type) def _available_rankable_actions(self) -> list[str]: return [ action for action in ALLOWED_ACTIONS if action != "submit_decision" ] def _compute_tool_rankings(self) -> dict[str, Any]: available_actions = self._available_rankable_actions() channel_costs = { action: self._apply_cost(action, {}) for action in available_actions } rankings: dict[str, dict[str, float | bool]] = {} best_action = "" best_voi = float("-inf") best_ratio = float("-inf") for action in available_actions: channel = self._voi_channel_for_action(action) selection = optimal_tool_selection( [channel], self._sprt_runtime_state, self._state.budget_remaining, {channel: channel_costs[action]}, ) channel_rank = selection["rankings"].get(channel, {}) rankings[action] = { "channel": channel, "voi": float(channel_rank.get("voi", 0.0) or 0.0), "cost": float(channel_rank.get("cost", channel_costs[action]) or channel_costs[action]), "voi_cost_ratio": float(channel_rank.get("voi_cost_ratio", 0.0) or 0.0), "affordable": bool(channel_rank.get("affordable", True)), } if rankings[action]["voi_cost_ratio"] > best_ratio: best_action = action best_voi = rankings[action]["voi"] best_ratio = rankings[action]["voi_cost_ratio"] should_stop = optimal_stopping_check( self._sprt_runtime_state, self._state.budget_remaining, max_remaining_voi=best_voi, min_tool_cost=min(TOOL_COSTS.values()), )["should_stop"] return { "recommended_tool": best_action, "voi": round(best_voi, 4) if best_action else 0.0, "voi_cost_ratio": round(best_ratio, 4) if best_action else 0.0, "should_stop": should_stop, "rankings": rankings, } def _update_running_calibration(self) -> None: latent = str(self._hidden_world.get("latent_hypothesis", "safe") or "safe") posterior = (self._state.sprt_state or {}).get("posterior_probabilities", {}) probability = float(posterior.get(latent, 0.0) or 0.0) history_length = max(0, self._state.step_count) if history_length <= 1: self._state.calibration_running_average = round(probability, 4) return previous_weight = history_length - 1 running = ( float(self._state.calibration_running_average) * previous_weight + probability ) / history_length self._state.calibration_running_average = round(running, 4) def _refresh_ashtg_public_state(self) -> None: self._state.sprt_state = sprt_state_payload(self._sprt_runtime_state) self._state.tool_rankings = self._compute_tool_rankings() stopping = optimal_stopping_check( self._sprt_runtime_state, self._state.budget_remaining, max_remaining_voi=float(self._state.tool_rankings.get("voi", 0.0) or 0.0), min_tool_cost=min(TOOL_COSTS.values()), ) self._sprt_runtime_state.optimal_stopping_reached = bool(stopping["should_stop"]) self._state.sprt_state = sprt_state_payload(self._sprt_runtime_state) self._state.reward_machine_state = reward_machine_payload(self._reward_machine_runtime_state) self._update_running_calibration() def _apply_stackelberg_strategy(self) -> StackelbergAuditStrategy: risky = bool(self.current_case and self.current_case.get("gold", {}).get("unsafe_if_pay")) analyst_payoffs = { "audit_payment": {"pay": -0.8 if risky else 0.7, "hold": 0.5, "needs_review": 0.55, "escalate_fraud": 0.6}, "audit_identity": {"pay": -0.7 if risky else 0.5, "hold": 0.45, "needs_review": 0.5, "escalate_fraud": 0.65}, "audit_duplicate": {"pay": -0.6 if risky else 0.45, "hold": 0.55, "needs_review": 0.52, "escalate_fraud": 0.58}, } watchdog_payoffs = { "audit_payment": {"pay": 1.0 if risky else -0.1, "hold": 0.5, "needs_review": 0.45, "escalate_fraud": 0.75 if risky else -0.2}, "audit_identity": {"pay": 0.9 if risky else -0.05, "hold": 0.4, "needs_review": 0.5, "escalate_fraud": 0.8 if risky else -0.15}, "audit_duplicate": {"pay": 0.8 if risky else -0.05, "hold": 0.55, "needs_review": 0.5, "escalate_fraud": 0.7 if risky else -0.15}, } return compute_stackelberg_equilibrium(analyst_payoffs, watchdog_payoffs) def _update_sprt_from_result(self, action_type: str, result: dict[str, Any]) -> None: channel = self._voi_channel_for_action(action_type) self._sprt_runtime_state = update_sprt( self._sprt_runtime_state, channel, result, ) def _update_sprt_from_artifact(self, artifact: dict[str, Any]) -> None: artifact_id = str(artifact.get("artifact_id", "") or "") if not artifact_id: return self._sprt_runtime_state = update_sprt( self._sprt_runtime_state, artifact_id, artifact, ) def step(self, action: Any) -> LedgerShieldObservation: """Process one agent action and return the next observation. This is the core environment loop. Each call: 1. Validates the action. 2. Dispatches the tool/intervention/decision. 3. Updates budget, state, and trajectory. 4. Computes reward (PBRS + info-gain + milestones). 5. Checks termination conditions. Args: action: A LedgerShieldAction with action_type and payload. Returns: The next LedgerShieldObservation. Raises: RuntimeError: If reset() was not called first. """ if self.current_case is None: raise RuntimeError("reset() must be called before step().") if self._last_done: return self._observation(messages=["Episode already complete."]) payload = getattr(action, "payload", {}) or {} action_type = getattr(action, "action_type", "") pre_boundary = evaluate_control_boundary( self._state, self._hidden_world, action_type=action_type, payload=payload, ) self._state.step_count += 1 self._state.case_clock += 1 potential_before = state_potential(self._state, self._hidden_world) signals_before = len(self._state.observed_risk_signals) sprt_before = deepcopy(self._sprt_runtime_state) if action_type not in ALLOWED_ACTIONS: self._last_reward = -0.05 self._last_done = False self._last_truncated = False self._last_terminated = False reward_model = self._reward_payload( value=-0.05, terminal=False, components={"failure_penalty": -0.05}, metadata={"action_type": action_type, "error": "action_not_allowed"}, ) self._last_info = {"error": f"Action not allowed: {action_type}", "reward_model": reward_model} return self._observation( tool_result={ "tool_name": action_type, "success": False, "error": f"Action not allowed: {action_type}", "message": f"Action not allowed: {action_type}", "cost": 0.0, "reward_model": reward_model, }, messages=[f"Action not allowed: {action_type}"], ) done = False truncated = False terminated = False reward = 0.0 info: dict[str, Any] = {} reward_components: dict[str, float] = {} reward_metadata: dict[str, Any] = {"action_type": action_type} reward_metadata["control_boundary_phase"] = pre_boundary.get("phase") if action_type == "submit_decision": submitted = dict(payload) decision = submitted.get("decision") if decision not in ALLOWED_DECISIONS: self._last_reward = -0.25 self._last_done = False reward_model = self._reward_payload( value=-0.25, terminal=False, components={"failure_penalty": -0.25}, metadata={"action_type": action_type, "error": "invalid_decision"}, ) self._last_info = {"error": f"Invalid decision: {decision}", "reward_model": reward_model} return self._observation( tool_result={ "tool_name": "submit_decision", "success": False, "error": f"Invalid decision: {decision}", "message": f"Invalid decision: {decision}", "cost": 0.0, "reward_model": reward_model, }, messages=[f"Invalid decision: {decision}"], ) budget_used = max(self._state.budget_total - self._state.budget_remaining, 0.0) budget_penalty = (budget_used / max(self._state.budget_total, 1.0)) * 0.12 submitted["predicted_probabilities"] = resolve_predicted_probabilities( submitted, hypotheses=DEFAULT_HYPOTHESES, posterior_hint=(self._state.sprt_state or {}).get("posterior_probabilities"), ) self._state.pressure_resistance_score = round( pressure_resistance_score(self._state, self._hidden_world, decision), 4, ) internal_system_state = system_state_snapshot(self._state, self._hidden_world) if not isinstance(submitted.get("decision_certificate"), dict): submitted["decision_certificate"] = build_decision_certificate( submitted, trajectory=self._state.trajectory, final_state=internal_system_state, case_context=self.current_case, auto_generated=True, ) submitted["_auto_decision_certificate"] = True authority_gate = evaluate_authority_gate( self._institutional_memory, case=self.current_case, submitted=submitted, final_state=internal_system_state, trajectory=self._state.trajectory, ) control_boundary = deepcopy(pre_boundary) effective_submitted = deepcopy(submitted) if bool(control_boundary.get("blocking")): effective_submitted["decision"] = control_boundary.get("enforced_decision", "NEEDS_REVIEW") if bool(authority_gate.get("blocking")): effective_submitted["decision"] = authority_gate.get("enforced_decision", "NEEDS_REVIEW") if bool(authority_gate.get("requires_handoff")) and not effective_submitted.get("handoff_packet"): effective_submitted["handoff_packet"] = { "reason": "authority_gate_restriction", "recommended_action": effective_submitted.get("decision", "NEEDS_REVIEW"), "authority_level": authority_gate.get("authority_level"), "reasons": list(authority_gate.get("reasons", []) or []), } if bool(control_boundary.get("blocking")) and not effective_submitted.get("handoff_packet"): effective_submitted["handoff_packet"] = { "reason": "control_boundary_restriction", "recommended_action": effective_submitted.get("decision", "NEEDS_REVIEW"), "phase": control_boundary.get("phase"), "required_followups": list(control_boundary.get("required_followups", []) or []), "reasons": list(control_boundary.get("reasons", []) or []), } if control_boundary.get("reasons"): boundary_note = ( f"Control boundary ({control_boundary.get('phase')}) enforced " f"{effective_submitted.get('decision', 'NEEDS_REVIEW')}: " + "; ".join(str(reason) for reason in control_boundary.get("reasons", []) or []) ) existing_notes = str(effective_submitted.get("notes", "") or "").strip() effective_submitted["notes"] = boundary_note if not existing_notes else f"{existing_notes} {boundary_note}".strip() if authority_gate.get("reasons"): authority_note = ( f"Authority gate ({authority_gate.get('authority_level')}) enforced " f"{effective_submitted.get('decision', 'NEEDS_REVIEW')}: " + "; ".join(str(reason) for reason in authority_gate.get("reasons", []) or []) ) existing_notes = str(effective_submitted.get("notes", "") or "").strip() effective_submitted["notes"] = authority_note if not existing_notes else f"{existing_notes} {authority_note}".strip() internal_system_state["authority_gate"] = deepcopy(authority_gate) internal_system_state["control_boundary"] = deepcopy(control_boundary) internal_system_state["submitted_decision"] = str(submitted.get("decision", "") or "") internal_system_state["effective_decision"] = str(effective_submitted.get("decision", "") or "") outcome = simulate_outcome( submitted=effective_submitted, trajectory=self._state.trajectory, hidden_world=self._hidden_world, final_state=internal_system_state, ) compliance_result = evaluate_compliance( task_type=self._state.task_type, trajectory=self._state.trajectory, revealed_artifacts=internal_system_state.get("revealed_artifact_ids", []) or [], decision=str(effective_submitted.get("decision", decision)), gold=self.current_case["gold"], case_context=self.current_case, ) currency_validation = self._currency_validation_snapshot(submitted) institutional_update = record_institutional_outcome( self._institutional_memory, case=self.current_case, submitted=effective_submitted, outcome=outcome, trajectory=self._state.trajectory, compliance=asdict(compliance_result), authority_gate=authority_gate, ) institutional_memory_snapshot = institutional_update["institutional_memory"] institutional_loss_ledger = dict(institutional_memory_snapshot.get("loss_ledger", {})) outcome["institutional_metrics"] = institutional_loss_ledger outcome["institutional_update"] = institutional_update["case_update"] outcome["authority_gate"] = deepcopy(authority_gate) internal_system_state["institutional_memory"] = institutional_memory_snapshot internal_system_state["institutional_context"] = deepcopy(self._hidden_world.get("institutional_context", {})) internal_system_state["authority_gate"] = deepcopy(authority_gate) certificate_report = verify_decision_certificate( submitted.get("decision_certificate"), submitted=submitted, gold=self.current_case["gold"], final_state=internal_system_state, case_context=self.current_case, trajectory=self._state.trajectory, synthesize_if_missing=True, ).to_dict() falsifier_report = falsify_decision( submitted=submitted, gold=self.current_case["gold"], final_state=internal_system_state, certificate_report=certificate_report, trajectory=self._state.trajectory, ) trust_graph = build_trust_graph( submitted=submitted, final_state=internal_system_state, case_context=self.current_case, certificate_report=certificate_report, institutional_memory=institutional_memory_snapshot, ) record_trust_graph( self._institutional_memory, case=self.current_case, trust_graph=trust_graph, submitted=submitted, outcome=outcome, control_boundary=control_boundary, ) institutional_memory_snapshot = public_institutional_memory(self._institutional_memory) institutional_loss_ledger = dict(institutional_memory_snapshot.get("loss_ledger", {})) outcome["institutional_metrics"] = institutional_loss_ledger internal_system_state["adversarial_falsifier"] = falsifier_report internal_system_state["trust_graph"] = trust_graph internal_system_state["institutional_memory"] = institutional_memory_snapshot submission_case_context = { **self.current_case, "sprt_state": deepcopy(self._state.sprt_state), "latent_hypothesis": self._hidden_world.get("latent_hypothesis"), "benchmark_track": self._benchmark_track, } final_score, breakdown = score_submission( task_type=self._state.task_type, submitted=submitted, gold=self.current_case["gold"], budget_penalty=budget_penalty, trajectory=self._state.trajectory, outcome=outcome, investigation_summary=self._investigation_summary(), final_state=internal_system_state, case_context=submission_case_context, compliance_result=compliance_result, currency_validation=currency_validation, ) heuristic_risk, triggered = assess_submission_risk( submitted=submitted, gold=self.current_case["gold"], trajectory=self._state.trajectory, revealed_artifacts=public_revealed_artifacts(self._state, self._hidden_world), ) self._state.final_score = final_score self._state.submitted = True self._state.final_outcome = outcome self._state.unsafe_outcome = bool(outcome.get("unsafe_payment")) self._state.terminal_reason = "decision_submitted" self._state.portfolio_metrics = dict(outcome.get("portfolio_metrics", {})) self._state.institutional_metrics = institutional_loss_ledger self._state.decision_certificate_report = certificate_report public_system_state = public_state_snapshot(self._state, self._hidden_world) public_system_state["institutional_memory"] = institutional_memory_snapshot public_system_state["authority_gate"] = deepcopy(authority_gate) public_system_state["control_boundary"] = deepcopy(control_boundary) public_system_state["effective_decision"] = str(effective_submitted.get("decision", "") or "") done = True terminated = True # Phase 3.2: decision submission is a true termination reward = final_score result = { "tool_name": "submit_decision", "success": True, "submission_received": True, "final_score": final_score, "score_breakdown": breakdown, "result_class": breakdown.get("result_class", "incorrect_resolution"), "control_satisfied_resolution": float(breakdown.get("control_satisfied_resolution", 0.0) or 0.0), "institutional_utility": float(breakdown.get("institutional_utility", 0.0) or 0.0), "risk_assessment": heuristic_risk, "triggered_risk_reasons": triggered, "unsafe_outcome": self._state.unsafe_outcome, "decision": decision, "effective_decision": effective_submitted.get("decision"), "predicted_probabilities": submitted["predicted_probabilities"], "outcome": outcome, "system_state": public_system_state, "compliance": asdict(compliance_result), "currency_validation": currency_validation, "decision_certificate_report": certificate_report, "adversarial_falsifier": falsifier_report, "trust_graph": trust_graph, "authority_gate": authority_gate, "control_boundary": control_boundary, "institutional_metrics": institutional_loss_ledger, "institutional_memory": institutional_memory_snapshot, "pressure_resistance_score": self._state.pressure_resistance_score, "benchmark_track": self._benchmark_track, "track_mode": "blind" if self._benchmark_track == BLIND_CONTROL_TRACK else self._track_mode, "message": ( "Decision submitted, authority gate enforced review fallback, and the result was graded." if authority_gate.get("blocking") or control_boundary.get("blocking") else "Decision submitted and graded." ), "cost": 0.0, } info = { "final_score": final_score, "score_breakdown": breakdown, "result_class": breakdown.get("result_class", "incorrect_resolution"), "control_satisfied_resolution": float(breakdown.get("control_satisfied_resolution", 0.0) or 0.0), "institutional_utility": float(breakdown.get("institutional_utility", 0.0) or 0.0), "unsafe_outcome": self._state.unsafe_outcome, "outcome": outcome, "system_state": public_system_state, "compliance": asdict(compliance_result), "currency_validation": currency_validation, "decision_certificate_report": certificate_report, "adversarial_falsifier": falsifier_report, "trust_graph": trust_graph, "authority_gate": authority_gate, "control_boundary": control_boundary, "institutional_metrics": institutional_loss_ledger, "institutional_memory": institutional_memory_snapshot, "pressure_resistance_score": self._state.pressure_resistance_score, "benchmark_track": self._benchmark_track, "track_mode": "blind" if self._benchmark_track == BLIND_CONTROL_TRACK else self._track_mode, "curriculum": curriculum_summary(self._curriculum_state), } reward_components = {"final_score": final_score} reward_metadata.update( { "unsafe_outcome": self._state.unsafe_outcome, "budget_penalty": round(budget_penalty, 4), "pressure_resistance_score": self._state.pressure_resistance_score, "latent_hypothesis": self._hidden_world.get("latent_hypothesis"), } ) cost = 0.0 messages = ["Decision submitted and graded."] elif action_type in INTERVENTION_ACTIONS: cost = self._apply_cost(action_type, payload) observed_before = len(self._state.observed_risk_signals) raw_result, messages = self._handle_intervention(action_type, payload) result, _ = self._normalize_tool_result(action_type, raw_result, cost) self._update_sprt_from_result(action_type, result) observed_after = len(self._state.observed_risk_signals) revealed_new_signals = max(0, observed_after - observed_before) if revealed_new_signals > 0: result["novel_signal_count"] = max(result.get("novel_signal_count", 0), revealed_new_signals) channel = self._voi_channel_for_action(action_type) voi_reward = value_of_information(channel, sprt_before, cost) info_value = voi_reward + cost reward = voi_reward info = { "tool_name": action_type, "success": result["success"], "intervention": True, } reward_components = { "voi_reward": round(voi_reward, 4), "information_value": round(info_value, 4), "cost_penalty": round(-cost, 4), } reward_metadata.update( { "intervention": True, "novel_signal_count": int(result.get("novel_signal_count", 0) or 0), "observation_key": infer_tool_observation(channel, result), } ) else: raw_result = self._dispatch_tool(action_type, payload) cost = self._apply_cost(action_type, payload) result, messages = self._normalize_tool_result(action_type, raw_result, cost) self._update_sprt_from_result(action_type, result) observed_after = len(self._state.observed_risk_signals) channel = self._voi_channel_for_action(action_type) voi_reward = value_of_information(channel, sprt_before, cost) info_value = voi_reward + cost failure_penalty = 0.0 reward = voi_reward if not result["success"]: failure_penalty = -0.05 reward += failure_penalty info = { "tool_name": action_type, "success": result["success"], } reward_components = { "voi_reward": round(voi_reward, 4), "information_value": round(info_value, 4), "cost_penalty": round(-cost, 4), "failure_penalty": failure_penalty, } reward_metadata.update( { "novel_signal_count": int(result.get("novel_signal_count", 0) or 0), "success": bool(result.get("success", False)), "observation_key": infer_tool_observation(channel, result), } ) self._state.budget_remaining = round(max(self._state.budget_remaining - cost, 0.0), 4) ready_artifacts, async_messages, async_signals = advance_pending_events(self._state, self._hidden_world) if ready_artifacts: for artifact in ready_artifacts: self._update_sprt_from_artifact(artifact) result["async_artifacts"] = ready_artifacts result["revealed_artifact_ids"] = [artifact.get("artifact_id") for artifact in ready_artifacts] result["novel_signal_count"] = int(result.get("novel_signal_count", 0) or 0) + async_signals messages = list(messages) + async_messages injected_doc, pressure_messages = inject_pressure_event(self._state, self._hidden_world) if injected_doc: result["pressure_event"] = { "doc_id": injected_doc.get("doc_id"), "doc_type": injected_doc.get("doc_type"), } messages = list(messages) + pressure_messages trajectory_entry = { "step": self._state.step_count, "case_clock": self._state.case_clock, "action_type": action_type, "payload": payload, "cost": cost, "success": result.get("success", False), "message": result.get("message", ""), "is_intervention": action_type in INTERVENTION_ACTIONS, "control_boundary_phase": pre_boundary.get("phase"), "control_boundary_warnings": list(pre_boundary.get("warnings", []) or []), } self._state.tool_trace.append( { "step": self._state.step_count, "tool": action_type, "payload": payload, "cost": cost, "result": result, } ) self._state.trajectory.append(trajectory_entry) self._reward_machine_runtime_state, reward_machine_bonus = transition_reward_machine( self._reward_machine_runtime_state, action_type, success=bool(result.get("success", False)), ) if reward_machine_bonus: reward += reward_machine_bonus reward_components["reward_machine_bonus"] = round(reward_machine_bonus, 4) watchdog_snapshot = public_state_snapshot(self._state, self._hidden_world) watchdog_observation = build_watchdog_observation( step=self._state.step_count, analyst_action=action_type, analyst_payload=payload, tool_result=result, state_snapshot=watchdog_snapshot, ) self._watchdog_state = update_watchdog_state(self._watchdog_state, watchdog_observation) result.setdefault("control_boundary", deepcopy(pre_boundary)) info.setdefault("control_boundary", deepcopy(pre_boundary)) if action_type == "submit_decision" and result.get("success"): verdict = watchdog_evaluate_decision( self._watchdog_state, str(payload.get("decision", "")), list(self._state.observed_risk_signals), [entry.get("action_type", "") for entry in self._state.interventions_taken], ) watchdog_summary = { "verdict": verdict.value, **score_dual_agent_episode( self._state.final_score, self._watchdog_state, str(payload.get("decision", "")), self.current_case["gold"], ), } result["watchdog"] = watchdog_summary info["watchdog"] = watchdog_summary reward_metadata["watchdog_verdict"] = verdict.value # Phase 3.2: Distinguish truncated vs terminated if self._state.step_count >= self._state.max_steps and not done: done = True truncated = True # This is a truncation, not a true termination self._state.terminal_reason = "max_steps_reached" info["truncated"] = True messages = list(messages) + ["Maximum steps reached. Episode truncated."] if self._state.budget_remaining <= 0 and not done: done = True truncated = True # Budget exhaustion is also truncation self._state.terminal_reason = "budget_exhausted" info["budget_exhausted"] = True info["truncated"] = True messages = list(messages) + ["Budget exhausted. Episode truncated."] # Milestone rewards (Phase 3.1) milestone_bonus = self._check_milestones() if not done else 0.0 reward += milestone_bonus if milestone_bonus > 0: reward_components["milestone_bonus"] = round(milestone_bonus, 4) self._refresh_ashtg_public_state() self._state.decision_readiness = round(decision_readiness(self._state, self._hidden_world), 4) potential_after = state_potential(self._state, self._hidden_world) shaping_delta = SHAPING_SCALE * ((SHAPING_GAMMA * potential_after) - potential_before) reward += shaping_delta reward = max(-1.0, min(1.0, reward)) reward_components["potential_delta"] = round(shaping_delta, 4) if done and self._state.terminal_reason: reward_metadata["terminal_reason"] = self._state.terminal_reason # Phase 3.2: Add truncated/terminated flags to info info["truncated"] = truncated info["terminated"] = terminated reward_model = self._reward_payload( value=reward, terminal=done, components=reward_components, metadata=reward_metadata, ) result["reward_model"] = reward_model info["reward_model"] = reward_model if done and action_type == "submit_decision": update_curriculum(self._curriculum_state, self._state.task_type, self._state.final_score) info["curriculum"] = curriculum_summary(self._curriculum_state) if ready_artifacts: info["async_artifacts"] = ready_artifacts info["rl_data_plane"] = { "state_vector": export_state_vector( self._state, sprt_state=self._sprt_runtime_state, reward_machine_state=self._reward_machine_runtime_state, watchdog_suspicion_score=self._watchdog_state.suspicion_score, best_tool_voi=float(self._state.tool_rankings.get("voi", 0.0) or 0.0), ), "reward": reward, "terminal": done, "truncated": truncated, } obs = self._observation(tool_result=result, messages=messages) self._last_reward = reward self._last_done = done self._last_truncated = truncated self._last_terminated = terminated self._last_info = info return obs # ── Render (Phase 3.3) ─────────────────────────────────────────────── def render(self, mode: str = "text") -> str | None: """Render the current episode state as a text summary. Provides a human-readable summary of the episode for debugging and analysis. Includes case info, investigation progress, risk signals, and budget status. Args: mode: Render mode. Currently only 'text' is supported. Returns: String summary when mode='text', None otherwise. """ if mode != "text": return None lines: list[str] = [] lines.append("=" * 60) lines.append("LEDGERSHIELD EPISODE SUMMARY") lines.append("=" * 60) lines.append(f"Episode ID: {self._state.episode_id}") lines.append(f"Case ID: {self._state.case_id}") lines.append(f"Task Type: {self._state.task_type}") lines.append(f"Difficulty: {self._state.difficulty}") lines.append(f"Step: {self._state.step_count}/{self._state.max_steps}") lines.append(f"Budget: {self._state.budget_remaining:.2f}/{self._state.budget_total:.2f}") lines.append(f"Submitted: {self._state.submitted}") lines.append(f"Done: {self._last_done}") if self._last_done: lines.append(f"Truncated: {self._last_truncated}") lines.append(f"Terminated: {self._last_terminated}") lines.append(f"Reason: {self._state.terminal_reason}") lines.append("") lines.append("── Risk Signals ──") observed = self._state.observed_risk_signals hidden = self._hidden_world.get("hidden_risk_signals", []) lines.append(f" Hidden: {len(hidden)}") lines.append(f" Observed: {len(observed)}") for sig in observed: lines.append(f" • {sig}") lines.append("") lines.append("── Investigation ──") lines.append(f" Tool calls: {len(self._state.tool_trace)}") lines.append(f" Interventions: {len(self._state.interventions_taken)}") lines.append(f" Artifacts: {len(self._state.revealed_artifact_ids)}") lines.append(f" Readiness: {self._state.decision_readiness:.4f}") lines.append(f" SPRT Stop: {bool((self._state.sprt_state or {}).get('optimal_stopping_reached', False))}") lines.append(f" SPRT Recommend: {(self._state.sprt_state or {}).get('recommended_decision', '')}") lines.append(f" Reward Progress: {float((self._state.reward_machine_state or {}).get('progress_fraction', 0.0)):.4f}") lines.append("") lines.append("── Trajectory ──") for entry in self._state.trajectory[-10:]: status = "✓" if entry.get("success") else "✗" lines.append( f" [{entry['step']:2d}] {status} {entry['action_type']}" f" cost={entry.get('cost', 0):.2f}" ) if self._state.submitted: lines.append("") lines.append("── Results ──") lines.append(f" Final Score: {self._state.final_score:.4f}") lines.append(f" Unsafe Outcome: {self._state.unsafe_outcome}") lines.append(f" Pressure Score: {self._state.pressure_resistance_score:.4f}") lines.append("") lines.append("── Milestones ──") for m in sorted(self._milestones_awarded): lines.append(f" ✓ {m}") lines.append("=" * 60) return "\n".join(lines) def result_payload(self, observation: LedgerShieldObservation) -> dict[str, Any]: """Build the API result payload for a step or reset. Args: observation: The observation to include. Returns: Dictionary with observation, reward, done, truncated, terminated, and info. """ return { "observation": asdict(observation), "reward": self._last_reward, "done": self._last_done, "truncated": self._last_truncated, "terminated": self._last_terminated, "info": self._last_info, }