| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import re |
| from collections import Counter |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| from openai import OpenAI |
|
|
| from client import SepsisTreatmentEnv |
| from models import SepsisAction, SepsisObservation |
|
|
|
|
| OUTPUT_DIR = Path("outputs") |
| MAX_STEPS_PER_TASK = {"easy": 8, "medium": 12, "hard": 16} |
| EPSILON = 0.1 |
| RNG = random.Random(7) |
| ENV_NAME = "sepsis-openenv" |
| LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") |
| DEFAULT_API_BASE_URL = "https://router.huggingface.co/v1" |
| DEFAULT_MODEL_NAME = "Qwen/Qwen2.5-72B-Instruct" |
| TASK_IDS = ["easy", "medium", "hard"] |
| LAB_OPTIONS = ["lactate", "wbc", "creatinine", "bicarbonate", "platelets", "bilirubin"] |
| TREATMENT_OPTIONS = ["monitor", "fluids", "vasopressors", "combination"] |
| LAB_ALIASES = { |
| "lactate": ["lactate", "lactic acid", "serum lactate", "blood lactate"], |
| "wbc": ["wbc", "white blood cell", "white blood cell count", "complete blood count", "cbc"], |
| "creatinine": [ |
| "creatinine", |
| "renal panel", |
| "kidney function", |
| "bmp", |
| "basic metabolic panel", |
| "cmp", |
| "comprehensive metabolic panel", |
| "bun", |
| ], |
| "bicarbonate": ["bicarbonate", "hco3", "co2", "carbon dioxide", "blood gas", "abg", "vbg"], |
| "platelets": ["platelets", "platelet count"], |
| "bilirubin": ["bilirubin", "total bili", "total bilirubin", "liver function", "lft"], |
| } |
| TREATMENT_ALIASES = { |
| "monitor": ["monitor", "observe", "observation", "watch", "watchful waiting", "reassess"], |
| "fluids": ["fluids", "iv fluids", "fluid resuscitation", "crystalloid", "bolus"], |
| "vasopressors": ["vasopressors", "pressor", "pressors", "norepinephrine", "levophed"], |
| "combination": ["combination", "fluids and vasopressors", "dual therapy", "both fluids and pressors"], |
| } |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="Run sepsis environment inference and aggregate metrics.") |
| parser.add_argument( |
| "--episodes", |
| type=int, |
| default=1, |
| help="Number of evaluation cycles to run. Each cycle runs easy, medium, and hard once.", |
| ) |
| parser.add_argument( |
| "--model", |
| choices=["auto", "heuristic", "llm", "id3qne"], |
| default="auto", |
| help="Policy mode. auto uses llm if credentials are present, otherwise heuristic.", |
| ) |
| parser.add_argument( |
| "--output", |
| default=str(OUTPUT_DIR / "baseline_scores.json"), |
| help="Path for the JSON summary output.", |
| ) |
| return parser.parse_args() |
|
|
|
|
| def format_action(action: SepsisAction) -> str: |
| if action.action_type == "request_lab": |
| return f"request_lab({action.lab_type}, suspect_sepsis={str(action.suspect_sepsis).lower()})" |
| if action.action_type == "request_treatment": |
| return f"request_treatment({action.treatment_type}, suspect_sepsis={str(action.suspect_sepsis).lower()})" |
| return f"monitor(suspect_sepsis={str(action.suspect_sepsis).lower()})" |
|
|
|
|
| def curriculum_action(observation: SepsisObservation) -> SepsisAction | None: |
| task_id = observation.task_id |
| step_index = observation.step_index |
|
|
| if task_id == "easy": |
| if step_index == 0: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="lactate", |
| rationale="Curriculum schedule: start with lactate for early workup.", |
| ) |
| if step_index == 2: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="creatinine", |
| rationale="Curriculum schedule: add renal assessment when the workup broadens.", |
| ) |
| return SepsisAction( |
| action_type="monitor", |
| suspect_sepsis=True, |
| rationale="Curriculum schedule: maintain observation after priority labs.", |
| ) |
|
|
| if task_id == "medium": |
| lab_plan = {0: "lactate", 1: "wbc", 2: "creatinine", 3: "bicarbonate"} |
| if step_index in lab_plan: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type=lab_plan[step_index], |
| rationale=f"Curriculum schedule: collect {lab_plan[step_index]} before treatment.", |
| ) |
| treatment_type = "fluids" if step_index <= 6 else "monitor" |
| return SepsisAction( |
| action_type="request_treatment", |
| suspect_sepsis=True, |
| treatment_type=treatment_type, |
| rationale="Curriculum schedule: shift from early support to monitoring.", |
| ) |
|
|
| if task_id == "hard": |
| if step_index == 0: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="lactate", |
| rationale="Curriculum schedule: start unstable trajectory with lactate.", |
| ) |
| if step_index == 1: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="creatinine", |
| rationale="Curriculum schedule: check renal strain before extended management.", |
| ) |
| treatment_type = "monitor" if step_index in {3, 4} else "fluids" |
| return SepsisAction( |
| action_type="request_treatment", |
| suspect_sepsis=True, |
| treatment_type=treatment_type, |
| rationale="Curriculum schedule: alternate stabilization and support across the harder case.", |
| ) |
|
|
| return None |
|
|
|
|
| def heuristic_action(observation: SepsisObservation) -> SepsisAction: |
| scheduled_action = curriculum_action(observation) |
| if scheduled_action is not None: |
| return scheduled_action |
|
|
| severity = observation.severity_proxy |
| shock = observation.vitals.get("Shock_Index", 0.0) |
| mean_bp = observation.vitals.get("MeanBP", 0.0) |
| visible_labs = observation.visible_labs |
| requested_labs = set(observation.requested_labs) |
|
|
| if RNG.random() < EPSILON: |
| unseen_labs = [lab for lab in ["lactate", "wbc", "creatinine", "bicarbonate"] if lab not in requested_labs] |
| if unseen_labs: |
| lab_choice = unseen_labs[0] |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=severity >= 1.0 or shock > 0.1 or mean_bp < 0.0, |
| lab_type=lab_choice, |
| rationale="Exploration step", |
| ) |
|
|
| lab_priority_order = ["lactate", "wbc", "creatinine", "bicarbonate"] |
| for lab in lab_priority_order: |
| should_request = False |
| if lab == "lactate": |
| should_request = lab not in requested_labs |
| elif lab == "wbc": |
| should_request = lab not in requested_labs and (severity >= 0.75 or shock > 0.08) |
| elif lab == "creatinine": |
| should_request = lab not in requested_labs and severity >= 1.2 |
| elif lab == "bicarbonate": |
| should_request = lab not in requested_labs and (severity >= 1.5 or mean_bp < -0.1) |
|
|
| if should_request: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=severity >= 1.0 or shock > 0.1 or mean_bp < 0.0, |
| lab_type=lab, |
| rationale=f"Exploring informative lab: {lab}", |
| ) |
|
|
| lactate = visible_labs.get("lactate", 0.0) or 0.0 |
| bicarbonate = visible_labs.get("bicarbonate", 0.0) or 0.0 |
| if severity < 0.8 and mean_bp >= -0.1: |
| treatment_type = "monitor" |
| elif severity >= 2.0 or mean_bp < -0.2: |
| treatment_type = "combination" if lactate > 0.25 else "vasopressors" |
| elif shock > 0.15 or severity >= 1.1 or bicarbonate < -0.15: |
| treatment_type = "fluids" |
| else: |
| treatment_type = "monitor" |
|
|
| return SepsisAction( |
| action_type="request_treatment", |
| suspect_sepsis=severity >= 1.0 or lactate > 0.25, |
| treatment_type=treatment_type, |
| rationale="Improved staged policy with exploration and severity awareness.", |
| ) |
|
|
|
|
| def iter_text_fragments(value: Any) -> list[str]: |
| if value is None: |
| return [] |
| if isinstance(value, str): |
| return [value] |
| if isinstance(value, (list, tuple, set)): |
| fragments: list[str] = [] |
| for item in value: |
| fragments.extend(iter_text_fragments(item)) |
| return fragments |
| if isinstance(value, dict): |
| fragments: list[str] = [] |
| for item in value.values(): |
| fragments.extend(iter_text_fragments(item)) |
| return fragments |
| return [str(value)] |
|
|
|
|
| def normalize_text(value: Any) -> str: |
| fragments = iter_text_fragments(value) |
| raw = " ".join(fragment.strip().lower() for fragment in fragments if fragment) |
| return re.sub(r"[^a-z0-9]+", " ", raw).strip() |
|
|
|
|
| def match_alias(value: Any, alias_map: dict[str, list[str]]) -> str | None: |
| fragments = iter_text_fragments(value) |
| matches: list[str] = [] |
| for fragment in fragments: |
| normalized = normalize_text(fragment) |
| if not normalized: |
| continue |
| for canonical, aliases in alias_map.items(): |
| if normalized == canonical: |
| matches.append(canonical) |
| continue |
| if any(alias in normalized for alias in aliases): |
| matches.append(canonical) |
| continue |
|
|
| if not matches: |
| combined = normalize_text(value) |
| for canonical, aliases in alias_map.items(): |
| if combined == canonical or any(alias in combined for alias in aliases): |
| matches.append(canonical) |
|
|
| unique_matches = list(dict.fromkeys(matches)) |
| if len(unique_matches) == 1: |
| return unique_matches[0] |
| return None |
|
|
|
|
| def parse_boolish(value: Any, default: bool = False) -> bool: |
| if isinstance(value, bool): |
| return value |
| if value is None: |
| return default |
| normalized = normalize_text(value) |
| if normalized in {"true", "yes", "y", "1"}: |
| return True |
| if normalized in {"false", "no", "n", "0"}: |
| return False |
| return default |
|
|
|
|
| def normalize_lab_choice(value: Any) -> str | None: |
| return match_alias(value, LAB_ALIASES) |
|
|
|
|
| def normalize_treatment_choice(value: Any) -> str | None: |
| return match_alias(value, TREATMENT_ALIASES) |
|
|
|
|
| def normalize_action_type(value: Any, lab_choice: str | None, treatment_choice: str | None) -> str | None: |
| normalized = normalize_text(value) |
| if normalized in {"request lab", "lab", "labs", "test", "request tests", "request lab test"}: |
| return "request_lab" |
| if normalized in {"request treatment", "treatment", "treat", "therapy", "intervene"}: |
| return "request_treatment" |
| if normalized in {"monitor", "observe", "observation", "watch", "reassess"}: |
| return "monitor" |
| if lab_choice: |
| return "request_lab" |
| if treatment_choice: |
| return "request_treatment" |
| return None |
|
|
|
|
| def should_use_heuristic_guardrail( |
| candidate: SepsisAction, |
| heuristic: SepsisAction, |
| observation: SepsisObservation, |
| ) -> bool: |
| del observation |
| return ( |
| candidate.action_type != heuristic.action_type |
| or candidate.lab_type != heuristic.lab_type |
| or candidate.treatment_type != heuristic.treatment_type |
| ) |
|
|
|
|
| def repair_model_action(payload: dict[str, Any], observation: SepsisObservation) -> tuple[SepsisAction, str, str | None]: |
| heuristic = heuristic_action(observation) |
| normalized_lab = normalize_lab_choice(payload.get("lab_type")) |
| normalized_treatment = normalize_treatment_choice(payload.get("treatment_type")) |
| normalized_action_type = normalize_action_type( |
| payload.get("action_type"), |
| normalized_lab, |
| normalized_treatment, |
| ) |
| suspect_sepsis = parse_boolish(payload.get("suspect_sepsis"), default=heuristic.suspect_sepsis) |
| suspect_sepsis = suspect_sepsis or heuristic.suspect_sepsis |
| rationale = str(payload.get("rationale", "")).strip() |
|
|
| if normalized_action_type == "request_lab": |
| if normalized_lab is None: |
| return heuristic, "heuristic_guardrail", "LLM selected an unsupported lab; using heuristic action." |
| candidate = SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=suspect_sepsis, |
| lab_type=normalized_lab, |
| rationale=rationale or "LLM lab choice normalized into environment action space.", |
| ) |
| elif normalized_action_type == "request_treatment": |
| if normalized_treatment is None: |
| return heuristic, "heuristic_guardrail", "LLM selected an unsupported treatment; using heuristic action." |
| candidate = SepsisAction( |
| action_type="request_treatment", |
| suspect_sepsis=suspect_sepsis, |
| treatment_type=normalized_treatment, |
| rationale=rationale or "LLM treatment choice normalized into environment action space.", |
| ) |
| elif normalized_action_type == "monitor": |
| candidate = SepsisAction( |
| action_type="monitor", |
| suspect_sepsis=suspect_sepsis, |
| rationale=rationale or "LLM monitor choice normalized into environment action space.", |
| ) |
| else: |
| return heuristic, "heuristic_guardrail", "LLM action could not be normalized; using heuristic action." |
|
|
| if should_use_heuristic_guardrail(candidate, heuristic, observation): |
| return heuristic, "heuristic_guardrail", "LLM action was valid but low-value for this step; using heuristic." |
|
|
| if candidate.model_dump(exclude={"rationale"}) == heuristic.model_dump(exclude={"rationale"}): |
| return candidate, "llm_aligned", None |
| return candidate, "llm_repaired", None |
|
|
|
|
| def id3qne_action(observation: SepsisObservation) -> SepsisAction: |
| task_id = observation.task_id |
| step_index = observation.step_index |
| severity = observation.severity_proxy |
| mean_bp = observation.vitals.get("MeanBP", 0.0) |
| shock = observation.vitals.get("Shock_Index", 0.0) |
| requested_labs = set(observation.requested_labs) |
| visible_labs = observation.visible_labs |
| suspect_sepsis = severity >= 1.0 or shock > 0.1 or mean_bp < 0.0 |
|
|
| if task_id == "easy": |
| if "lactate" not in requested_labs: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="lactate", |
| rationale="ID3QNE tree: always reveal lactate first in the easy workup branch.", |
| ) |
| if step_index >= 2 and "creatinine" not in requested_labs: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="creatinine", |
| rationale="ID3QNE tree: second split requests creatinine for renal assessment.", |
| ) |
| return SepsisAction( |
| action_type="monitor", |
| suspect_sepsis=True, |
| rationale="ID3QNE tree: monitor after the high-yield easy branch labs are collected.", |
| ) |
|
|
| if task_id == "medium": |
| for lab_name in ["lactate", "wbc", "creatinine", "bicarbonate"]: |
| if lab_name not in requested_labs: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type=lab_name, |
| rationale=f"ID3QNE tree: continue the medium-depth lab branch with {lab_name}.", |
| ) |
| treatment_type = "fluids" if step_index <= 6 else "monitor" |
| return SepsisAction( |
| action_type="request_treatment", |
| suspect_sepsis=True, |
| treatment_type=treatment_type, |
| rationale="ID3QNE tree: treat early, then monitor after stabilization.", |
| ) |
|
|
| if task_id == "hard": |
| if "lactate" not in requested_labs: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="lactate", |
| rationale="ID3QNE tree: unstable branch starts with lactate.", |
| ) |
| if "creatinine" not in requested_labs: |
| return SepsisAction( |
| action_type="request_lab", |
| suspect_sepsis=True, |
| lab_type="creatinine", |
| rationale="ID3QNE tree: follow with creatinine when the trajectory turns unstable.", |
| ) |
| creatinine = visible_labs.get("creatinine", 0.0) or 0.0 |
| if step_index in {3, 4} and severity < 1.5 and mean_bp >= -0.2: |
| treatment_type = "monitor" |
| elif severity >= 2.0 and mean_bp < -0.2: |
| treatment_type = "combination" |
| elif severity >= 1.0 or creatinine > 0.15 or step_index >= 5: |
| treatment_type = "fluids" |
| else: |
| treatment_type = "monitor" |
| return SepsisAction( |
| action_type="request_treatment", |
| suspect_sepsis=suspect_sepsis or creatinine > 0.15 or step_index >= 1, |
| treatment_type=treatment_type, |
| rationale="ID3QNE tree: treatment branch uses severity, renal strain, and step progression.", |
| ) |
|
|
| return heuristic_action(observation) |
|
|
|
|
| def build_prompt(observation: SepsisObservation) -> str: |
| return ( |
| "You are controlling a sequential sepsis management simulator.\n" |
| f"Task: {observation.task_description}\n" |
| f"Step: {observation.step_index + 1}/{observation.max_steps}\n" |
| f"Severity proxy: {observation.severity_proxy:.2f}\n" |
| f"Mortality flag in logged trajectory: {observation.mortality_risk_flag}\n" |
| f"Demographics: {json.dumps(observation.demographics)}\n" |
| f"Vitals: {json.dumps(observation.vitals)}\n" |
| f"Context features: {json.dumps(observation.context_features)}\n" |
| f"Visible labs: {json.dumps(observation.visible_labs)}\n" |
| f"Requested labs so far: {json.dumps(observation.requested_labs)}\n" |
| "You must choose exactly one environment action.\n" |
| f"Allowed lab_type values: {json.dumps(LAB_OPTIONS)}.\n" |
| f"Allowed treatment_type values: {json.dumps(TREATMENT_OPTIONS)}.\n" |
| "If action_type is request_lab, lab_type must be one of the allowed values and treatment_type must be null.\n" |
| "If action_type is request_treatment, treatment_type must be one of the allowed values and lab_type must be null.\n" |
| "If action_type is monitor, both lab_type and treatment_type must be null.\n" |
| "Do not return lists, synonyms, antibiotics, blood cultures, or free-text clinical plans.\n" |
| "Return JSON only with keys action_type, suspect_sepsis, lab_type, treatment_type, rationale." |
| ) |
|
|
|
|
| def parse_model_json(content: str) -> dict[str, Any]: |
| candidate = content.strip() |
| if candidate.startswith("```"): |
| lines = candidate.splitlines() |
| if lines and lines[0].startswith("```"): |
| lines = lines[1:] |
| if lines and lines[-1].startswith("```"): |
| lines = lines[:-1] |
| candidate = "\n".join(lines).strip() |
|
|
| match = re.search(r"\{.*\}", candidate, re.DOTALL) |
| if match: |
| candidate = match.group(0) |
| return json.loads(candidate) |
|
|
|
|
| def model_action( |
| client: OpenAI | None, |
| model_name: str | None, |
| observation: SepsisObservation, |
| ) -> tuple[SepsisAction, str, str | None]: |
| if client is None or not model_name: |
| return heuristic_action(observation), "heuristic", None |
|
|
| messages = [ |
| {"role": "system", "content": "Return only valid JSON for a sepsis management action."}, |
| {"role": "user", "content": build_prompt(observation)}, |
| ] |
| try: |
| completion = client.chat.completions.create( |
| model=model_name, |
| messages=messages, |
| temperature=0.0, |
| max_tokens=200, |
| response_format={"type": "json_object"}, |
| ) |
| content = completion.choices[0].message.content or "" |
| payload = parse_model_json(content) |
| return repair_model_action(payload, observation) |
| except Exception as exc: |
| return heuristic_action(observation), "heuristic_fallback", str(exc) |
|
|
|
|
| def choose_action( |
| policy_mode: str, |
| client: OpenAI | None, |
| model_name: str | None, |
| observation: SepsisObservation, |
| ) -> tuple[SepsisAction, str, str | None]: |
| if policy_mode == "heuristic": |
| return heuristic_action(observation), "heuristic", None |
| if policy_mode == "id3qne": |
| return id3qne_action(observation), "id3qne", None |
| if policy_mode == "llm": |
| return model_action(client, model_name, observation) |
| raise ValueError(f"Unsupported policy mode: {policy_mode}") |
|
|
|
|
| def compute_action_entropy(action_history: list[str]) -> float: |
| if not action_history: |
| return 0.0 |
|
|
| action_lengths = [len(action.split()) for action in action_history] |
| counts = np.bincount(action_lengths) |
| nonzero_counts = counts[counts > 0] |
| probabilities = nonzero_counts / len(action_history) |
| entropy = float(-np.sum(probabilities * np.log2(probabilities))) |
| if abs(entropy) < 1e-12: |
| return 0.0 |
| return entropy if entropy > 0 else 0.0 |
|
|
|
|
| def compute_dense_reward_metrics( |
| reward_trace: list[float], |
| step_count: int, |
| max_steps: int, |
| action_history: list[str], |
| ) -> dict[str, float | int]: |
| nonzero_rewards = [reward for reward in reward_trace if reward != 0] |
|
|
| return { |
| "steps_taken": step_count, |
| "total_reward": float(sum(reward_trace)), |
| "reward_count": len(reward_trace), |
| "positive_rewards_count": sum(1 for reward in reward_trace if reward > 0), |
| "reward_density": float(sum(1 for reward in reward_trace if reward > 0) / len(reward_trace)) |
| if reward_trace |
| else 0.0, |
| "avg_reward_per_step": float(np.mean(reward_trace)) if reward_trace else 0.0, |
| "reward_variance": float(np.var(reward_trace)) if reward_trace else 0.0, |
| "max_single_reward": float(max(reward_trace)) if reward_trace else 0.0, |
| "episode_length_efficiency": float(step_count / max_steps) if max_steps else 0.0, |
| "positive_reward_ratio": float( |
| sum(1 for reward in reward_trace if reward > 0) / max(1, len(nonzero_rewards)) |
| ), |
| "unique_actions": len(set(action_history)), |
| "action_entropy": compute_action_entropy(action_history), |
| } |
|
|
|
|
| def run_task( |
| task_id: str, |
| policy_mode: str, |
| client: OpenAI | None, |
| model_name: str | None, |
| episode_index: int, |
| ) -> dict[str, Any]: |
| global EPSILON |
| if task_id == "easy": |
| EPSILON = 0.05 |
| elif task_id == "medium": |
| EPSILON = 0.1 |
| else: |
| EPSILON = 0.15 |
|
|
| env = SepsisTreatmentEnv(base_url=os.getenv("ENV_BASE_URL"), task_id=task_id) |
| result = env.reset() |
| observation = result.observation |
| final_info = result.info |
| reward_trace: list[float] = [] |
| action_history: list[str] = [] |
| policy_sources: Counter[str] = Counter() |
| policy_errors: list[str] = [] |
| success = False |
| step_count = 0 |
|
|
| print( |
| f"[START] task={task_id} env={ENV_NAME} policy={policy_mode} " |
| f"model={model_name or policy_mode}" |
| ) |
|
|
| try: |
| for step_number in range(1, MAX_STEPS_PER_TASK[task_id] + 1): |
| action, source, error_message = choose_action(policy_mode, client, model_name, observation) |
| formatted_action = format_action(action) |
| result = env.step(action) |
| observation = result.observation |
| final_info = result.info |
| reward = float(result.reward or 0.0) |
| reward_trace.append(reward) |
| action_history.append(formatted_action) |
| policy_sources[source] += 1 |
| if error_message: |
| policy_errors.append(error_message) |
| step_count = step_number |
| print( |
| f"[STEP] step={step_number} source={source} action={formatted_action} " |
| f"reward={reward:.2f} done={str(result.done).lower()} error=null" |
| ) |
| if result.done: |
| success = True |
| break |
| except Exception as exc: |
| policy_errors.append(str(exc)) |
| success = False |
| finally: |
| state = env.state() |
| env.close() |
| rewards_repr = ",".join(f"{reward:.2f}" for reward in reward_trace) |
| print(f"[END] success={str(success).lower()} steps={step_count} rewards={rewards_repr}") |
|
|
| metrics = final_info.get("metrics", {}) |
| dense_metrics = compute_dense_reward_metrics( |
| reward_trace=reward_trace, |
| step_count=step_count, |
| max_steps=MAX_STEPS_PER_TASK[task_id], |
| action_history=action_history, |
| ) |
| return { |
| "task_id": task_id, |
| "episode_id": state.episode_id, |
| "score": metrics.get("score", 0.0), |
| "avg_reward": metrics.get("avg_reward", 0.0), |
| "detection": metrics.get("detection", 0.0), |
| "lab_workup": metrics.get("lab_workup", 0.0), |
| "treatment": metrics.get("treatment", 0.0), |
| "timeliness": metrics.get("timeliness", 0.0), |
| "stability": metrics.get("stability", 0.0), |
| "safety": metrics.get("safety", 0.0), |
| "safety_violation_rate": metrics.get("safety_violation_rate", 0.0), |
| "safety_violations": metrics.get("safety_violations", 0), |
| "outcome": metrics.get("outcome", 0.0), |
| "steps": metrics.get("steps", state.step_count), |
| "episode_index": episode_index, |
| "policy_mode": policy_mode, |
| "policy_sources": dict(policy_sources), |
| "policy_error_count": len(policy_errors), |
| "policy_last_error": policy_errors[-1] if policy_errors else None, |
| **dense_metrics, |
| } |
|
|
|
|
| def summarize_runs( |
| all_results: list[dict[str, Any]], |
| per_episode_results: list[dict[str, Any]], |
| requested_policy: str, |
| active_policy: str, |
| model_name: str, |
| ) -> dict[str, Any]: |
| if not all_results: |
| raise ValueError("No results were generated.") |
|
|
| policy_source_totals: Counter[str] = Counter() |
| for result in all_results: |
| policy_source_totals.update(result.get("policy_sources", {})) |
|
|
| total_reward_count = sum(result["reward_count"] for result in all_results) |
| total_positive_rewards = sum(result["positive_rewards_count"] for result in all_results) |
| total_steps = sum(result["steps_taken"] for result in all_results) |
| total_safety_violations = sum(result["safety_violations"] for result in all_results) |
|
|
| return { |
| "results": all_results, |
| "episode_summaries": per_episode_results, |
| "mean_score": round(float(np.mean([item["score"] for item in all_results])), 4), |
| "score_std": round(float(np.std([item["score"] for item in all_results])), 4), |
| "mean_score_std": round(float(np.std([item["mean_score"] for item in per_episode_results])), 4) |
| if per_episode_results |
| else 0.0, |
| "mean_reward_density": round(float(np.mean([item["reward_density"] for item in all_results])), 4), |
| "global_reward_density": round(float(total_positive_rewards / total_reward_count), 4) |
| if total_reward_count |
| else 0.0, |
| "mean_avg_reward_per_step": round(float(np.mean([item["avg_reward_per_step"] for item in all_results])), 4), |
| "mean_reward_variance": round(float(np.mean([item["reward_variance"] for item in all_results])), 4), |
| "mean_positive_reward_ratio": round(float(np.mean([item["positive_reward_ratio"] for item in all_results])), 4), |
| "mean_action_entropy": round(float(np.mean([item["action_entropy"] for item in all_results])), 4), |
| "safety_violation_rate": round(float(total_safety_violations / total_steps), 4) if total_steps else 0.0, |
| "total_runs": len(all_results), |
| "episodes": len(per_episode_results), |
| "requested_policy": requested_policy, |
| "active_policy": active_policy, |
| "model_name": model_name, |
| "policy_source_totals": dict(policy_source_totals), |
| } |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| OUTPUT_DIR.mkdir(exist_ok=True) |
| api_base_url = os.getenv("API_BASE_URL", DEFAULT_API_BASE_URL) |
| model_name = os.getenv("MODEL_NAME", DEFAULT_MODEL_NAME) |
| api_key = os.getenv("OPENAI_API_KEY") or os.getenv("HF_TOKEN") |
|
|
| llm_client = None |
| if api_base_url and model_name and api_key: |
| llm_client = OpenAI(base_url=api_base_url, api_key=api_key) |
|
|
| if args.episodes < 1: |
| raise SystemExit("--episodes must be at least 1.") |
|
|
| if args.model == "llm" and llm_client is None: |
| raise SystemExit("LLM mode requires OPENAI_API_KEY or HF_TOKEN plus API_BASE_URL and MODEL_NAME.") |
|
|
| active_policy = args.model |
| if args.model == "auto": |
| active_policy = "llm" if llm_client is not None else "heuristic" |
|
|
| all_results: list[dict[str, Any]] = [] |
| episode_summaries: list[dict[str, Any]] = [] |
| for episode_index in range(args.episodes): |
| episode_results = [ |
| run_task(task_id, active_policy, llm_client, model_name, episode_index) for task_id in TASK_IDS |
| ] |
| all_results.extend(episode_results) |
| episode_steps = sum(item["steps_taken"] for item in episode_results) |
| episode_safety_violations = sum(item["safety_violations"] for item in episode_results) |
| episode_summaries.append( |
| { |
| "episode_index": episode_index, |
| "mean_score": round(float(np.mean([item["score"] for item in episode_results])), 4), |
| "mean_reward_density": round(float(np.mean([item["reward_density"] for item in episode_results])), 4), |
| "safety_violation_rate": round(float(episode_safety_violations / episode_steps), 4) |
| if episode_steps |
| else 0.0, |
| } |
| ) |
|
|
| summary = summarize_runs( |
| all_results=all_results, |
| per_episode_results=episode_summaries, |
| requested_policy=args.model, |
| active_policy=active_policy, |
| model_name=model_name if active_policy == "llm" else active_policy, |
| ) |
| output_path = Path(args.output) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| output_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
