"""Baseline inference script for the Hospital ED OpenEnv environment. Hackathon requirements fulfilled by this file: * File is named ``inference.py`` and lives in the repo root. * Uses the OpenAI client for all LLM calls. * Reads credentials from environment variables: - ``API_BASE_URL`` — LLM endpoint (default: HF router) - ``MODEL_NAME`` — model identifier (default: Qwen2.5-72B-Instruct) - ``HF_TOKEN`` / ``API_KEY`` — bearer token - ``LOCAL_IMAGE_NAME`` / ``IMAGE_NAME`` — Docker image name for :meth:`HospitalEdEnv.from_docker_image` * Runs 3 tasks with a difficulty progression (easy → medium → hard): 1. ``normal_day`` — steady-state ED operations (easy) 2. ``surge`` — COVID-style respiratory surge (medium) 3. ``mass_casualty`` — trauma burst, rapid triage (hard) * Emits the exact ``[START] / [STEP] / [END]`` stdout format specified in the hackathon "Mandatory Additional Instructions" section. * Produces per-task scores in [0, 1] (``task_score`` from the env metadata), plus a combined aggregate. Runtime budget: well under 20 minutes on vCPU=2 / 8 GB memory. Each task runs for at most 100 steps with a short system prompt and a fixed ``max_tokens=32`` response limit. """ from __future__ import annotations import asyncio import os import re import textwrap import traceback from typing import List, Optional from openai import OpenAI from client import HospitalEdEnv from models import HospitalAction # --------------------------------------------------------------------- # Environment variables # --------------------------------------------------------------------- API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1" MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct" API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or "dummy" IMAGE_NAME = ( os.getenv("LOCAL_IMAGE_NAME") or os.getenv("IMAGE_NAME") or "hospital-ed-env:latest" ) BENCHMARK = os.getenv("HOSPITAL_ED_BENCHMARK", "hospital_ed") # Runtime caps — keeps total inference time well under 20 minutes even # with a slow LLM. 100 env steps × 3 tasks × ~2 s/LLM call ≈ 10 min. MAX_STEPS_PER_TASK = 100 TEMPERATURE = 0.2 MAX_TOKENS = 32 # we only need a single integer # Tasks in ascending difficulty, as required by the hackathon rubric. TASKS = ["normal_day", "surge", "mass_casualty"] SYSTEM_PROMPT = textwrap.dedent( """ You are the on-duty triage agent for a hospital emergency department. You control 20 general beds, 5 ICU beds, and 3 ventilators. Each timestep you pick exactly one discrete action from this table: 0 No-op (wait) 1-10 Assign waiting-queue patient [idx] to a GENERAL bed 11-20 Assign waiting-queue patient [idx] to an ICU bed 21-23 Use ventilator slot [idx] on the most severe ICU patient who does not already have one 24-28 Transfer general bed [idx] patient to an ICU bed 29-33 Discharge general bed [idx] patient 34-38 Discharge ICU bed [idx] patient Severity scale: 1 = mild, 2 = moderate, 3 = severe, 4 = critical. Critical (4) patients die if they wait too long — prioritise them into the ICU first. Attach ventilators to severe respiratory cases. You will be given the current hospital state and a list of valid actions (a subset of 0-38). Reply with ONLY a single integer drawn from that valid list. No words, no punctuation, just the number. """ ).strip() # --------------------------------------------------------------------- # Logging helpers — exact stdout format from the hackathon spec. # --------------------------------------------------------------------- def log_start(task: str, env: str, model: str) -> None: print(f"[START] task={task} env={env} model={model}", flush=True) def log_step( step: int, action: str, reward: float, done: bool, error: Optional[str], ) -> None: error_val = error if error else "null" print( f"[STEP] step={step} action={action} reward={reward:.2f} " f"done={str(done).lower()} error={error_val}", flush=True, ) def log_end( success: bool, steps: int, score: float, rewards: List[float] ) -> None: rewards_str = ",".join(f"{r:.2f}" for r in rewards) print( f"[END] success={str(success).lower()} steps={steps} " f"score={score:.3f} rewards={rewards_str}", flush=True, ) # --------------------------------------------------------------------- # LLM policy # --------------------------------------------------------------------- def _format_observation(obs, valid_actions: List[int]) -> str: """Format a HospitalObservation into a compact text prompt.""" # Show only queue rows that have a waiting patient (severity > 0). queue_lines = [] for i, row in enumerate(obs.waiting_queue): sev = int(row[0]) if len(row) > 0 else 0 if sev == 0: continue cond = int(row[1]) if len(row) > 1 else 0 wait = int(row[2]) if len(row) > 2 else 0 queue_lines.append(f" queue[{i}] sev={sev} cond={cond} wait={wait}") queue_block = "\n".join(queue_lines) if queue_lines else " (empty)" stats = obs.stats or [0, 0, 0] treated = stats[0] if len(stats) > 0 else 0 deaths = stats[1] if len(stats) > 1 else 0 waiting = stats[2] if len(stats) > 2 else 0 return textwrap.dedent( f""" Hospital state at timestep {obs.time_step}: General beds (severity, 0=empty): {list(obs.bed_occupancy)} ICU beds (severity, 0=empty): {list(obs.icu_occupancy)} Ventilators (0=free, 1=in use): {list(obs.ventilator_status)} Waiting queue ({waiting} patients): {queue_block} Cumulative: treated={treated}, deaths={deaths} Valid actions right now: {valid_actions} Reply with ONE integer from the valid list. """ ).strip() def _parse_action( text: str, valid_actions: List[int], fallback: int = 0 ) -> int: """Parse the LLM response into an action id, clamped to valid_actions.""" match = re.search(r"-?\d+", text or "") if match is None: return fallback try: a = int(match.group(0)) except ValueError: return fallback if a in valid_actions: return a return fallback if fallback in valid_actions else valid_actions[0] def _ask_llm( client: OpenAI, obs, valid_actions: List[int] ) -> int: """Query the LLM for the next action and return a valid integer id.""" user_prompt = _format_observation(obs, valid_actions) try: completion = client.chat.completions.create( model=MODEL_NAME, messages=[ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": user_prompt}, ], temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False, ) text = (completion.choices[0].message.content or "").strip() except Exception as exc: print(f"[DEBUG] LLM request failed: {exc}", flush=True) text = "" fallback = valid_actions[0] if valid_actions else 0 return _parse_action(text, valid_actions, fallback=fallback) # --------------------------------------------------------------------- # Episode runner # --------------------------------------------------------------------- async def run_task( env: HospitalEdEnv, llm: OpenAI, task: str, seed: int, ) -> tuple[float, int, List[float], bool]: """Run one episode for ``task`` and return (score, steps, rewards, success).""" log_start(task=task, env=BENCHMARK, model=MODEL_NAME) rewards: List[float] = [] steps_taken = 0 score = 0.0 success = False last_error: Optional[str] = None obs = None try: result = await env.reset(seed=seed, task=task) obs = result.observation for step in range(1, MAX_STEPS_PER_TASK + 1): if result.done: break valid_actions = [i for i, m in enumerate(obs.action_mask) if m] if not valid_actions: valid_actions = [0] action_id = _ask_llm(llm, obs, valid_actions) try: result = await env.step(HospitalAction(action=action_id)) obs = result.observation reward = float(result.reward or 0.0) done = bool(result.done) last_error = None except Exception as step_exc: # pragma: no cover reward = 0.0 done = True last_error = str(step_exc) rewards.append(reward) steps_taken = step log_step( step=step, action=f"discrete({action_id})", reward=reward, done=done, error=last_error, ) if done: break # Final normalized task score is a direct Pydantic field on the # observation (metadata is stripped by OpenEnv's serializer). score = float(getattr(obs, "task_score", 0.0) or 0.0) score = max(0.0, min(1.0, score)) success = score >= 0.5 except Exception as exc: # pragma: no cover last_error = f"{type(exc).__name__}: {exc}" traceback.print_exc() finally: log_end( success=success, steps=steps_taken, score=score, rewards=rewards ) return score, steps_taken, rewards, success # --------------------------------------------------------------------- # Main # --------------------------------------------------------------------- async def main() -> None: llm = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) # Spin up the OpenEnv Hospital ED server in a Docker container and # connect to it. The hackathon grader expects exactly this pattern. env = await HospitalEdEnv.from_docker_image(IMAGE_NAME) try: per_task_scores = [] for i, task in enumerate(TASKS): score, _, _, _ = await run_task(env, llm, task=task, seed=42 + i) per_task_scores.append((task, score)) finally: try: await env.close() except Exception as e: # pragma: no cover print(f"[DEBUG] env.close() error: {e}", flush=True) # Aggregate summary printed after all [END] lines so judges can see # the composite score at a glance. if per_task_scores: mean_score = sum(s for _, s in per_task_scores) / len(per_task_scores) per_task_str = ", ".join( f"{name}={score:.3f}" for name, score in per_task_scores ) print( f"[SUMMARY] tasks={len(per_task_scores)} " f"mean_score={mean_score:.3f} {per_task_str}", flush=True, ) if __name__ == "__main__": asyncio.run(main())