SamarJamal/disaster_response_rl
1
1"""2Inference Script — AI Disaster Response Coordinator3===================================4MANDATORY5- Before submitting, ensure the following variables are defined in your environment configuration:6 HF_TOKEN Your Hugging Face API token (never commit the real value).7 API_BASE_URL The API endpoint for the LLM (OpenAI-compatible Hugging Face router).8 MODEL_NAME The model identifier to use for inference.9 IMAGE_NAME The name of the local image to use for the environment if you are using from_docker_image()10 11- Defaults are set for API_BASE_URL and MODEL_NAME:12 API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")13 MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")14 15- The inference script must be named `inference.py` and placed in the root directory of the project16- Participants must use OpenAI Client for all LLM calls using above variables17 18STDOUT FORMAT19- The script must emit exactly three line types to stdout, in this order:20 21 [START] task=<task_name> env=<benchmark> model=<model_name>22 [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>23 [END] success=<true|false> steps=<n> rewards=<r1,r2,...,rn>24 25 Rules:26 - One [START] line at episode begin.27 - One [STEP] line per step, immediately after env.step() returns.28 - One [END] line after env.close(), always emitted (even on exception).29 - reward and rewards are formatted to 2 decimal places.30 - done and success are lowercase booleans: true or false.31 - error is the raw last_action_error string, or null if none.32 - All fields on a single line with no newlines within a line.33 34 Example:35 [START] task=disaster_medium env=disaster_response model=Qwen/Qwen2.5-72B-Instruct36 [STEP] step=1 action={"assignments":[{"vehicle_id":"veh_1","location_id":"loc_1"}]} reward=51.00 done=false error=null37 [END] success=true steps=5 rewards=51.00,21.00,14.00,12.00,8.0038"""39 40import asyncio41import json42import os43import textwrap44from typing import Any, Dict, List, Optional45 46from openai import OpenAI47 48try:49 from models import DisasterAction, VehicleAssignment50 from client import DisasterResponseClient51except (ImportError, ModuleNotFoundError):52 from my_env import DisasterAction, DisasterResponseClient53 from my_env.models import VehicleAssignment54 55# ---------------------------------------------------------------------------56# Configuration — Hugging Face router via OpenAI-compatible client57# ---------------------------------------------------------------------------58 59HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")60if not HF_TOKEN:61 raise ValueError(62 "Set HF_TOKEN or OPENAI_API_KEY to your API token "63 "(create one at https://huggingface.co/settings/tokens — do not commit it)."64 )65 66API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")67MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")68 69# Docker image for environment container70IMAGE_NAME = os.getenv("IMAGE_NAME", "disaster_response-env:latest")71 72# Benchmark metadata73BENCHMARK = "disaster_response"74 75# Task configurations: name -> (difficulty, max_steps)76TASKS = {77 "disaster_easy": ("easy", 10),78 "disaster_medium": ("medium", 15),79 "disaster_hard": ("hard", 20),80}81 82# LLM parameters83TEMPERATURE = 0.3 # low temperature for deterministic decisions84MAX_TOKENS = 30085 86# ---------------------------------------------------------------------------87# System Prompt — instructs the LLM how to act in this environment88# ---------------------------------------------------------------------------89 90SYSTEM_PROMPT = textwrap.dedent(91 """\92 You are an AI Disaster Response Coordinator. Your goal is to rescue people from multiple locations using limited vehicles.93 94 ENVIRONMENT DYNAMICS:95 - Each vehicle can be assigned to exactly one location per step.96 - Vehicles rescue min(capacity, people_waiting) at the target location.97 - Locations with people left waiting incur increasing "waiting time" penalties.98 - Rewards: Rescue (+2.0), High Severity (+3.0 bonus), Medium Severity (+1.5 bonus).99 - Penalties: Idle vehicles (-1.0), Wasted dispatch (-2.0), Neglect (-0.5).100 101 STRATEGY (Follow these priorities):102 1. CAPACITY MATCHING: Match vehicle capacity to the number of people waiting. Don't waste a high-capacity vehicle on a site with few people if a smaller vehicle can handle it.103 2. SEVERITY-FIRST: Prioritize HIGH severity first, then MEDIUM, then LOW.104 3. WAIT-TIME MITIGATION: If multiple sites have the same severity, prioritize the one with the higher 'waiting_time' to minimize penalties.105 4. NO IDLE VEHICLES: Always assign every available vehicle to a site that still has people waiting.106 5. NO WASTED TRIPS: Never send a vehicle to a location where people_waiting is 0.107 108 OUTPUT FORMAT:109 You must provide your reasoning briefly, followed by the assignments in a JSON block code fence.110 111 Example Output:112 Reasoning: Location loc_1 has 20 people and is high priority. Assigning veh_1 (cap 10).113 ```json114 [{"vehicle_id": "veh_1", "location_id": "loc_1"}]115 ```116 """117).strip()118 119 120# ---------------------------------------------------------------------------121# Logging helpers (mandatory stdout format)122# ---------------------------------------------------------------------------123 124 125def log_start(task: str, env: str, model: str) -> None:126 print(f"[START] task={task} env={env} model={model}", flush=True)127 128 129def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:130 error_val = error if error else "null"131 done_val = str(done).lower()132 print(133 f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",134 flush=True,135 )136 137 138def log_end(success: bool, steps: int, rewards: List[float], score: float) -> None:139 rewards_str = ",".join(f"{r:.2f}" for r in rewards)140 print(141 f"[END] success={str(success).lower()} steps={steps} rewards={rewards_str} score={score:.4f}",142 flush=True,143 )144 145 146# ---------------------------------------------------------------------------147# Observation formatter — makes the state readable for the LLM148# ---------------------------------------------------------------------------149 150 151def format_observation_for_llm(obs: Any) -> str:152 """Convert a DisasterObservation into a clear text prompt for the LLM."""153 lines = []154 lines.append(f"TIME STEP: {obs.time_step} / {obs.max_steps}")155 lines.append(f"TOTAL PEOPLE SAVED SO FAR: {obs.total_people_saved}")156 lines.append(f"PEOPLE SAVED THIS STEP: {obs.people_saved_this_step}")157 lines.append("")158 159 # Locations160 lines.append("DISASTER LOCATIONS:")161 for loc in obs.locations:162 lines.append(163 f" - {loc.name} (id={loc.id}): severity={loc.severity}, "164 f"people_waiting={loc.people_waiting}, waiting_time={loc.waiting_time}"165 )166 lines.append("")167 168 # Vehicles169 lines.append("RESCUE VEHICLES:")170 for v in obs.vehicles:171 lines.append(172 f" - {v.name} (id={v.id}): capacity={v.capacity}, busy={v.is_busy}"173 )174 lines.append("")175 176 # Reward breakdown (if available)177 if obs.reward_breakdown:178 lines.append("LAST STEP REWARD BREAKDOWN:")179 for key, val in obs.reward_breakdown.items():180 lines.append(f" {key}: {val:+.2f}")181 lines.append("")182 183 lines.append(184 "Assign each available vehicle to a location. "185 "Respond with ONLY a JSON array of assignments."186 )187 return "\n".join(lines)188 189 190# ---------------------------------------------------------------------------191# LLM interaction192# ---------------------------------------------------------------------------193 194 195import re196 197def parse_llm_assignments(text: str) -> List[Dict[str, str]]:198 """Parse the LLM response into a list of assignment dicts, supporting reasoning text."""199 # Try to find JSON inside markdown blocks first200 json_match = re.search(r"```json\s*(\[.*?\])\s*```", text, re.DOTALL)201 if not json_match:202 # Try finding any array-like structure203 json_match = re.search(r"(\[.*\])", text, re.DOTALL)204 205 if json_match:206 cleaned = json_match.group(1).strip()207 try:208 parsed = json.loads(cleaned)209 if isinstance(parsed, list):210 valid = []211 for item in parsed:212 if (213 isinstance(item, dict)214 and "vehicle_id" in item215 and "location_id" in item216 ):217 valid.append({218 "vehicle_id": str(item["vehicle_id"]),219 "location_id": str(item["location_id"]),220 })221 return valid222 except json.JSONDecodeError:223 pass224 225 return []226 227 228def get_model_assignments(229 client: OpenAI,230 obs_text: str,231 history: List[str],232) -> List[Dict[str, str]]:233 """Call the LLM to get vehicle assignments for this step."""234 messages = [{"role": "system", "content": SYSTEM_PROMPT}]235 236 # Include recent history for context (last 3 steps)237 if history:238 history_text = "\n".join(history[-3:])239 messages.append({"role": "user", "content": f"Previous steps:\n{history_text}"})240 messages.append({"role": "assistant", "content": "Understood, I'll use this context."})241 242 messages.append({"role": "user", "content": obs_text})243 244 try:245 completion = client.chat.completions.create(246 model=MODEL_NAME,247 messages=messages,248 temperature=TEMPERATURE,249 max_tokens=MAX_TOKENS,250 stream=False,251 )252 text = (completion.choices[0].message.content or "").strip()253 assignments = parse_llm_assignments(text)254 if assignments:255 return assignments256 print(f"[DEBUG] Could not parse LLM response: {text!r}", flush=True)257 return []258 except Exception as exc:259 print(f"[DEBUG] Model request failed: {exc}", flush=True)260 return []261 262 263# ---------------------------------------------------------------------------264# Grader — computes normalized score [0.0, 1.0] for each task265# ---------------------------------------------------------------------------266 267 268def compute_task_score(269 total_people_saved: int,270 total_people_initial: int,271 rewards: List[float],272 done_success: bool,273 steps_taken: int,274 max_steps: int,275) -> float:276 """277 Compute a deterministic, normalized task score between 0.0 and 1.0.278 279 Scoring formula:280 - 60% weight: percentage of people rescued281 - 20% weight: completion bonus (1.0 if all rescued, 0.0 if timeout)282 - 20% weight: efficiency bonus (steps remaining / max_steps)283 """284 # Rescue percentage (0.0 to 1.0)285 rescue_pct = (286 total_people_saved / total_people_initial287 if total_people_initial > 0288 else 0.0289 )290 291 # Completion bonus292 completion_bonus = 1.0 if done_success else 0.0293 294 # Efficiency bonus: how many steps were saved295 efficiency = (296 (max_steps - steps_taken) / max_steps297 if max_steps > 0 and done_success298 else 0.0299 )300 301 score = (0.6 * rescue_pct) + (0.2 * completion_bonus) + (0.2 * efficiency)302 # Clamp score to strictly (0, 1) range as required by validator303 return min(max(score, 0.01), 0.99)304 305 306# Total people per difficulty (for grader)307TOTAL_PEOPLE = {308 "easy": 20, # 15 + 5309 "medium": 53, # 20 + 15 + 10 + 8310 "hard": 86, # 25 + 18 + 12 + 15 + 10 + 6311}312 313 314# ---------------------------------------------------------------------------315# Run a single task episode316# ---------------------------------------------------------------------------317 318 319async def run_task(320 llm_client: OpenAI,321 env: DisasterResponseClient,322 task_name: str,323 difficulty: str,324 max_steps: int,325) -> float:326 """327 Run a single task episode and return the normalized score.328 """329 history: List[str] = []330 rewards: List[float] = []331 steps_taken = 0332 score = 0.0333 success = False334 total_people_initial = TOTAL_PEOPLE.get(difficulty, 53)335 336 log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)337 338 try:339 # Set difficulty via environment variable before reset340 os.environ["DISASTER_DIFFICULTY"] = difficulty341 result = await env.reset()342 obs = result.observation343 344 for step in range(1, max_steps + 1):345 if result.done:346 break347 348 # Format the observation for the LLM349 obs_text = format_observation_for_llm(obs)350 351 # Get assignments from the LLM352 raw_assignments = get_model_assignments(llm_client, obs_text, history)353 354 # Build the action355 action = DisasterAction(356 assignments=[357 VehicleAssignment(358 vehicle_id=a["vehicle_id"],359 location_id=a["location_id"],360 )361 for a in raw_assignments362 ]363 )364 365 # Step the environment366 result = await env.step(action)367 obs = result.observation368 369 reward = result.reward or 0.0370 done = result.done371 error = None372 373 rewards.append(reward)374 steps_taken = step375 376 # Log in the required format377 action_str = json.dumps(378 {"assignments": raw_assignments}, separators=(",", ":")379 )380 log_step(step=step, action=action_str, reward=reward, done=done, error=error)381 382 # Build history entry for LLM context383 history.append(384 f"Step {step}: assigned {len(raw_assignments)} vehicles -> "385 f"rescued {obs.people_saved_this_step} people, reward={reward:+.2f}"386 )387 388 if done:389 break390 391 # Determine if episode ended with all people rescued392 all_rescued = all(loc.people_waiting == 0 for loc in obs.locations)393 394 # Compute normalized score using the grader395 score = compute_task_score(396 total_people_saved=obs.total_people_saved,397 total_people_initial=total_people_initial,398 rewards=rewards,399 done_success=all_rescued,400 steps_taken=steps_taken,401 max_steps=max_steps,402 )403 success = score >= 0.5404 405 finally:406 log_end(success=success, steps=steps_taken, rewards=rewards, score=score)407 408 return score409 410 411# ---------------------------------------------------------------------------412# Main entry point — runs all 3 tasks413# ---------------------------------------------------------------------------414 415 416async def main() -> None:417 llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)418 419 # Use the existing, running server at localhost:8000420 async with DisasterResponseClient(base_url="http://localhost:8000") as env:421 print("Connected to Local Environment. Starting evaluation...")422 scores = {}423 for task_name, (difficulty, max_steps) in TASKS.items():424 score = await run_task(llm_client, env, task_name, difficulty, max_steps)425 scores[task_name] = score426 print(f"[DEBUG] {task_name} score: {score:.3f}", flush=True)427 428 # Print final summary429 avg_score = sum(scores.values()) / len(scores)430 # Ensure average is also strictly within (0, 1)431 clamped_avg = min(max(avg_score, 0.01), 0.99)432 433 print(f"\n[SUMMARY] Average score: {clamped_avg:.4f}", flush=True)434 for name, sc in scores.items():435 print(f" {name}: {sc:.4f}", flush=True)436 437 438if __name__ == "__main__":439 asyncio.run(main())