AdityParbat/disaster-response-env

1# inference.py2"""3Inference Script — DisasterResponseEnv (RL Version - Submission Compliant)4Final Build: Hardened Rewards + 72B Intelligence + Strict Logging5"""6 7import argparse8import asyncio9import copy10import json11import os12import re13import time14import textwrap15import math16import urllib.request17from typing import List, Optional, Any18 19from openai import OpenAI20from models import Action, StepResult, Observation, Dispatch21from tasks import TASKS22 23# ==============================================================================24# CHECKBOX 2 & 3: STRICT ENVIRONMENT VARIABLES (Hugging Face / LLM Config)25# ==============================================================================26API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")27MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")28HF_TOKEN = os.getenv("HF_TOKEN")29LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME") # Optional30 31# Fallback to args only if env vars aren't set by the platform32parser = argparse.ArgumentParser(description="Disaster Response RL Evaluator")33parser.add_argument("--task", type=str, default=None)34parser.add_argument("--no-history", action="store_true", help="Disable history in prompt")35args, _ = parser.parse_known_args()36 37BENCHMARK = "disaster_response_env"38TEMPERATURE = 0.039MAX_TOKENS = 250040 41# ==============================================================================42# CHECKBOX 5: STRICT STDOUT LOGGING FORMAT (Regex Friendly)43# ==============================================================================44def log_start(task: str, env: str, model: str) -> None:45    print(f"[START] task={task} env={env} model={model}", flush=True)46 47def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:48    error_val = error if error else "null"49    print(50        f"[STEP] step={step} action={action} reward={reward:.2f} "51        f"done={str(done).lower()} error={error_val}",52        flush=True53    )54 55def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:56    rewards_str = ",".join(f"{r:.2f}" for r in rewards)57    print(58        f"[END] success={str(success).lower()} steps={steps} "59        f"score={score:.3f} rewards={rewards_str}",60        flush=True61    )62 63# ==============================================================================64# ENVIRONMENT WRAPPER (MDP Logic)65# ==============================================================================66class HTTPEnvWrapper:67    def __init__(self, task_id: str):68        self.task_id = task_id69        self.base_url = "https://adityparbat-disaster-response-env.hf.space"70 71    async def reset(self) -> StepResult:72        url = f"{self.base_url}/reset"73        data = json.dumps({"task_id": self.task_id}).encode('utf-8')74        req = urllib.request.Request(url, data=data, method='POST', headers={'Content-Type': 'application/json'})75        with urllib.request.urlopen(req) as response:76            return StepResult(**json.loads(response.read().decode()))77 78    async def step(self, action: Action) -> StepResult:79        url = f"{self.base_url}/step"80        # Combine task_id with the action dispatches81        payload = {82            "task_id": self.task_id,83            "dispatches": [d.model_dump() for d in action.dispatches],84            "recalls": [r.model_dump() for r in getattr(action, "recalls", [])],85            "reasoning": action.reasoning86        }87        data = json.dumps(payload).encode('utf-8')88        req = urllib.request.Request(url, data=data, method='POST', headers={'Content-Type': 'application/json'})89        with urllib.request.urlopen(req) as response:90            return StepResult(**json.loads(response.read().decode()))91 92# ==============================================================================93# AGENT LOGIC (Reasoning & Sanitization)94# ==============================================================================95async def get_model_action(client: OpenAI, obs: Observation, history: List[str]) -> Action:96    """Prompt Engineering: Resource-to-Incident Mapping Logic."""97    98    # 1. Specialized Manifest: Inject Identity Locks directly into the metadata99    specialized_manifest = copy.deepcopy(obs.resources_manifest)100    identity_msg = ""101    # Hardcoded context for the top-tier 'citywide' constraints102    if "unit_delta_4" in specialized_manifest:103        specialized_manifest["unit_delta_4"] = "fire_truck (IDENTITY LOCKED/REQUIRED for INC-003)"104        identity_msg = "- CRITICAL: INC-003 requires unit_delta_4. Standard fire_trucks will fail."105 106    # 2. Tactical Noise Reduction: Focus the model ONLY on what's relevant107    required_types = set()108    for inc in obs.active_incidents:109        required_types.update(inc.requires)110    111    filtered_available = []112    for u in obs.available_units:113        u_type = specialized_manifest.get(u, "")114        if any(u_type.startswith(req) for req in required_types):115            filtered_available.append(u)116    if not filtered_available: filtered_available = obs.available_units117 118    # 3. Prompt Construction119    SYSTEM_PROMPT = "You are the City Emergency Dispatcher. Coordinate specialized units to resolve active disasters."120 121    # Resource Advisor Table122    resource_rows = []123    for u in filtered_available:124        status = obs.busy_units.get(u, "Available")125        resource_rows.append(f"| {u} | {specialized_manifest.get(u)} | {status} |")126    resource_table = "\n".join(resource_rows)127 128    # Incident Table129    incident_rows = []130    for inc in obs.active_incidents:131        incident_rows.append(f"| {inc.id} | {inc.type} | {inc.severity} | {', '.join(inc.requires)} |")132    incident_table = "\n".join(incident_rows)133 134    history_block = "\n".join(history[-3:]) if history and not args.no_history else "No previous actions (this is Step 1)."135 136    user_prompt = textwrap.dedent(f"""137    Step {obs.step}/{obs.max_steps}138    139    ### [INCIDENT STATUS]140    | ID | Type | Severity | Requirements |141    | :--- | :--- | :--- | :--- |142    {incident_table}143 144    ### [RESOURCE STATUS]145    | ID | Type | Status (Turns Busy) |146    | :--- | :--- | :--- |147    {resource_table}148 149    ### [TACTICAL ADVICE]150    {identity_msg}151    - Map units to their specific requirement types.152    - DO NOT send redundant units of the same type to the same incident unless required.153    - DO NOT send busy units.154 155    ### [ACTION HISTORY]156    {history_block}157 158    ### [TASK]159    1. ANALYZE: Which incidents are unresolved?160    2. MAP: Match available units to their required types.161    3. RESPOND: provide a JSON object with "dispatches" and "reasoning".162    163    Your Response MUST be JSON:164    {{165      "dispatches": [ {{"unit": "id", "incident_id": "id"}} ],166      "reasoning": "Confirming [unit] meets [incident] requirement and is AVAILABLE."167    }}168    """).strip()169 170    try:171        completion = client.chat.completions.create(172            model=MODEL_NAME,173            messages=[174                {"role": "system", "content": SYSTEM_PROMPT},175                {"role": "user",   "content": user_prompt}176            ],177            temperature=TEMPERATURE,178            max_tokens=MAX_TOKENS179        )180        raw = completion.choices[0].message.content or ""181        match = re.search(r'\{.*\}', raw, re.DOTALL)182        if match:183            return Action(**json.loads(match.group()))184    except Exception as e:185        print(f"[DEBUG] API/Parse Error: {e}")186    return Action(dispatches=[], reasoning="error_fallback")187 188def sanitize_action(action: Action, obs: Observation) -> Action:189    """Clean-Up Wrapper: Enforces basic physical consistency (deduplication/availability)."""190    sane_recalls = []191    for r in getattr(action, "recalls", []):192        if r.unit in obs.busy_units:193            sane_recalls.append(r)194            if r.unit not in obs.available_units:195                obs.available_units.append(r.unit)196    action.recalls = sane_recalls197 198    sane_dispatches = []199    used_units = set()200    for d in action.dispatches:201        if d.unit in used_units or d.unit not in obs.available_units: 202            continue203        sane_dispatches.append(d)204        used_units.add(d.unit)205    action.dispatches = sane_dispatches206    207    return action208 209# ==============================================================================210# EXECUTION LOOP: SUBMISSION COMPLIANT211# ==============================================================================212async def run_task(client: OpenAI, task_id: str):213    env = HTTPEnvWrapper(task_id)214    215    # Checkbox 5: Log START216    log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME)217    218    result = await env.reset()219    total_reward = 0.0220    history = []221    rewards = []222    steps_taken = 0223    success = False224    225    for step in range(1, result.observation.max_steps + 1):226        # 1. Prediction227        action = await get_model_action(client, result.observation, history)228        229        # 2. Cleanup (Common Sense Deduplication)230        action = sanitize_action(action, result.observation)231        232        # 3. Environment Step233        result = await env.step(action)234        235        # 4. State Management236        total_reward += result.reward237        rewards.append(result.reward)238        steps_taken = step239        done = result.terminated or result.truncated240        241        # Format action and error for the strict logger242        action_json = json.dumps([{"unit": d.unit, "incident_id": d.incident_id} for d in action.dispatches]).replace(" ", "")243        error_msg = ", ".join(result.info.get("violations", [])) if result.info.get("violations") else None244        245        # Checkbox 5: Log STEP246        log_step(step=step, action=action_json, reward=result.reward, done=done, error=error_msg)247        248        history.append(f"Step {step}: Action={action_json}, Reward={result.reward:.1f}")249        250        if done:251            success = result.terminated252            break253            254    # Checkbox 5: Logistic Normalization for Grader Compliance (0.01 - 0.99)255    # This maps our unbounded MDP return into the strictly bounded range required by the grader.256    sigmoid = 1 / (1 + math.exp(-total_reward))257    final_score = 0.01 + (0.98 * sigmoid)258    259    log_end(success=success, steps=steps_taken, score=final_score, rewards=rewards)260    return {"task_id": task_id, "reward": total_reward, "success": success}261 262async def main():263    if not HF_TOKEN:264        print("[ERROR] HF_TOKEN environment variable is missing. The automated checker requires it.")265        return266 267    # Checkbox 4: Client instantiation using required variables268    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)269    270    tasks = [args.task] if args.task else list(TASKS.keys())271    for t in tasks:272        await run_task(client, t)273 274if __name__ == "__main__":275    asyncio.run(main())276