PunitRaveendran/Air_Traffic_Control_System

1"""Inference script for ATC environment using OpenAI client."""2import json3import os4import re5from typing import List, Optional, Tuple6 7from openai import OpenAI8 9from env.atc_env import ATCEnv, Action10from env.aircraft import AircraftStatus, Priority11from graders.grader1 import grade as grade112from graders.grader2 import grade as grade213from graders.grader3 import grade as grade314 15# ==========================================16# HACKATHON COMPLIANT LOGGING FUNCTIONS17# ==========================================18def log_start(task: str, env: str, model: str) -> None:19    print(f"[START] task={task} env={env} model={model}", flush=True)20 21def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:22    error_val = error if error else "null"23    done_val = str(done).lower()24    print(25        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",26        flush=True,27    )28 29def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:30    rewards_str = ",".join(f"{r:.2f}" for r in rewards)31    clamped_score = min(max(score, 0.0), 1.0)32    print(f"[END] success={str(success).lower()} steps={steps} score={clamped_score:.3f} rewards={rewards_str}", flush=True)33# ==========================================34 35 36def will_exhaust_fuel(ac: dict) -> bool:37    """True if this aircraft will run out of fuel before it can reach the runway."""38    return ac.get("fuel_remaining_min", 999) < ac.get("eta_steps", 999)39 40 41def get_sort_key(ac: dict):42    """43    Urgency sort (ascending = handle first):44      0: fuel < ETA  → will crash without immediate action45      1: EMERGENCY priority46      2: FUEL_CRITICAL priority47      3: fuel < 15 (about to become fuel-critical)48      4: NORMAL49    Secondary: least fuel first, then earliest ETA.50    """51    fuel = ac.get("fuel_remaining_min", 999)52    eta = ac.get("eta_steps", 999)53    priority = ac.get("priority", "NORMAL")54 55    if will_exhaust_fuel(ac):56        tier = 057    elif priority == "EMERGENCY":58        tier = 159    elif priority == "FUEL_CRITICAL":60        tier = 261    elif fuel < 15:62        tier = 363    else:64        tier = 465 66    return (tier, fuel, eta)67 68 69def get_fallback_action(env: ATCEnv, task_id: int = 1, current_step: int = 0) -> List[Action]:70    """71    Deterministic greedy controller — used when LLM is bypassed or fails.72 73    Key behaviours:74    - Sorts all unassigned aircraft by urgency (crash risk > emergency > fuel_critical > normal).75    - Task 3, step >= 2: avoids RW02 (it closes at step 3; stranded assignments are now76      reset by the env on closure, but better not to assign there in the first place).77    - Assigns ALL unassigned aircraft every step to keep queue minimal.78    - Round-robin across available runways with monotonically increasing sequence positions.79    """80    state = env.state().model_dump()81    aircraft_list = state.get("aircraft", [])82 83    unassigned = [84        ac for ac in aircraft_list85        if ac.get("status") not in ("LANDED", "LANDING", "ASSIGNED")86    ]87 88    if not unassigned:89        return []90 91    ordered = sorted(unassigned, key=get_sort_key)92 93    runways = state.get("runways", [])94    available_runways = [rw for rw in runways if rw.get("status") != "CLOSED"]95 96    # Task 3: proactively avoid RW02 at step >= 2 since it closes at step 3.97    # The env will reset stuck aircraft on closure, but avoiding assignment is cleaner.98    if task_id == 3 and current_step >= 2:99        available_runways = [rw for rw in available_runways if rw["id"] != "RW02"]100 101    if not available_runways:102        return []103 104    runway_ids = [rw["id"] for rw in available_runways]105    runway_queue_pos = {rw["id"]: 1 for rw in available_runways}106 107    actions = []108    for i, ac in enumerate(ordered):109        runway_id = runway_ids[i % len(runway_ids)]110        seq_pos = runway_queue_pos[runway_id]111        runway_queue_pos[runway_id] += 1112 113        fuel = ac.get("fuel_remaining_min", 999)114        priority = ac.get("priority", "NORMAL")115        is_critical = (116            will_exhaust_fuel(ac)117            or priority in ("EMERGENCY", "FUEL_CRITICAL")118            or fuel <= 15119        )120        action_type = "expedite" if is_critical else "assign"121 122        actions.append(Action(123            aircraft_id=ac["id"],124            action_type=action_type,125            runway_id=runway_id,126            sequence_position=seq_pos,127        ))128 129    return actions130 131 132def should_skip_llm(task_id: int, current_step: int) -> bool:133    """134    Returns True to bypass the LLM and use deterministic controller instead.135 136    Task 2: Always skip — fuel-critical aircraft (AC3, AC6) must be expedited in step 1.137            LLM API failures or mis-ordering are catastrophic here.138    Task 3: Skip for steps 1-12 to reliably handle:139            - Steps 1-2: assign all 15 initial aircraft before RW02 closes140            - Step 3: RW02 closes (env auto-reassigns stuck aircraft to HOLDING)141            - Step 5: AC10 becomes EMERGENCY — must be expedited immediately142            - Steps 4-12: batch arrivals at steps 4, 8 need immediate assignment143    """144    if task_id == 2:145        return True146    if task_id == 3 and current_step <= 12:147        return True148    return False149 150 151def parse_llm_response(response_text: str) -> List[dict]:152    """Parse LLM response, stripping markdown fences before extracting JSON."""153    clean = re.sub(r'```(?:json)?', '', response_text).strip()154    clean = clean.replace('```', '')155 156    match = re.search(r'\{.*\}', clean, re.DOTALL)157    if match:158        try:159            data = json.loads(match.group(0))160            if isinstance(data, dict) and "actions" in data:161                return data["actions"]162            if isinstance(data, list):163                return data164        except json.JSONDecodeError:165            pass166 167    return []168 169 170def build_prompt(observation: dict, state: dict, task_id: int) -> str:171    """Task-specific prompts targeting each grader's exact scoring components."""172    aircraft_list = state.get("aircraft", [])173    current_step = state.get("timestep", 0)174 175    # Show only actionable aircraft, sorted by urgency176    actionable = [177        ac for ac in aircraft_list178        if ac.get("status") in ("INBOUND", "HOLDING")179    ]180    sorted_aircraft = sorted(actionable, key=get_sort_key)181 182    aircraft_lines = []183    for ac in sorted_aircraft:184        priority = ac.get("priority", "NORMAL")185        fuel = ac.get("fuel_remaining_min", 0)186        dist = ac.get("distance_nm", 0)187        eta = ac.get("eta_steps", 999)188        warning = ""189        if will_exhaust_fuel(ac):190            warning = " *** CRASH RISK: fuel < ETA — EXPEDITE NOW ***"191        elif priority in ("EMERGENCY", "FUEL_CRITICAL") or fuel < 15:192            warning = " *** CRITICAL ***"193        aircraft_lines.append(194            f"- {ac['id']}: priority={priority}, fuel={fuel:.1f}min, "195            f"dist={dist:.1f}nm, ETA={eta}steps{warning}"196        )197 198    runway_lines = []199    for rw in state.get("runways", []):200        status = rw.get("status", "UNKNOWN")201        note = " (DO NOT USE)" if status == "CLOSED" else ""202        runway_lines.append(f"- {rw['id']}: {status}{note}")203 204    if task_id == 1:205        task_focus = ("Maximize aircraft landed. Assign all INBOUND/HOLDING aircraft "206                      "to runways immediately. Spread load across runways.")207    elif task_id == 2:208        task_focus = (209            "SCORED: 40% landings + 40% fuel-critical lands + 20% no duplicate positions.\n"210            "CRITICAL: Aircraft marked CRASH RISK must be expedited first or the episode continues "211            "with a severe penalty. Always expedite FUEL_CRITICAL aircraft.\n"212            "Never give two aircraft the same sequence_position on the same runway."213        )214    else:215        extra = ""216        if current_step >= 2:217            extra += "\nWARNING: RW02 is closed or closing — assign ALL aircraft to RW01 only."218        if current_step >= 5:219            extra += "\nAC10 is EMERGENCY — expedite immediately if not yet LANDED."220        task_focus = (221            "SCORED: 25% landings + 25% emergency lands + 25% no fuel exhaustion + 25% queue<=12.\n"222            "Assign ALL unassigned aircraft every step. Expedite any CRASH RISK or EMERGENCY aircraft."223            + extra224        )225 226    prompt = f"""You are an expert ATC sequencing controller.227Timestep: {current_step}228 229Aircraft needing assignment (most urgent first):230{chr(10).join(aircraft_lines) if aircraft_lines else "  (none — all aircraft already assigned or landed)"}231 232Runways:233{chr(10).join(runway_lines)}234 235{task_focus}236 237RULES:2381. Never assign to a CLOSED runway.2392. Unique sequence_position per runway: if 3 aircraft go to RW01 use positions 1, 2, 3.2403. Use "expedite" for EMERGENCY, FUEL_CRITICAL, or CRASH RISK aircraft.2414. Use "assign" for NORMAL aircraft.2425. Only include INBOUND or HOLDING aircraft in actions.243 244Respond ONLY with this JSON:245```json246{{247  "thought_process": "brief reasoning",248  "actions": [249    {{"aircraft_id": "AC3", "action_type": "expedite", "runway_id": "RW01", "sequence_position": 1}},250    {{"aircraft_id": "AC1", "action_type": "assign", "runway_id": "RW01", "sequence_position": 2}}251  ]252}}253```"""254    return prompt255 256 257def apply_guardrails(action_dicts: List[dict], state_dump: dict) -> List[Action]:258    """259    Validate and sanitize LLM action dicts before converting to Action objects.260    Re-sorts by urgency and enforces unique sequence positions per runway.261    """262    open_runways = [263        rw["id"] for rw in state_dump.get("runways", [])264        if rw.get("status") != "CLOSED"265    ]266    aircraft_map = {a["id"]: a for a in state_dump.get("aircraft", [])}267 268    # Sort by urgency before assigning positions269    action_dicts_sorted = sorted(270        action_dicts,271        key=lambda ad: get_sort_key(aircraft_map.get(ad.get("aircraft_id", ""), {}))272    )273 274    runway_queue_pos = {}275    actions = []276 277    for action_dict in action_dicts_sorted:278        try:279            ac_id = action_dict.get("aircraft_id", "")280            action_type = action_dict.get("action_type", "hold")281            runway_id = action_dict.get("runway_id")282 283            ac_data = aircraft_map.get(ac_id)284            if ac_data:285                is_critical = (286                    will_exhaust_fuel(ac_data)287                    or ac_data.get("fuel_remaining_min", 999) <= 15288                    or ac_data.get("priority") in ("EMERGENCY", "FUEL_CRITICAL")289                )290                # Never hold a critical aircraft291                if action_type == "hold" and is_critical:292                    action_type = "expedite"293                # Critical aircraft must have an open runway294                if is_critical and (not runway_id or runway_id not in open_runways):295                    if open_runways:296                        runway_id = open_runways[0]297 298            # Never assign to closed runway299            if runway_id and runway_id not in open_runways:300                if open_runways:301                    runway_id = open_runways[0]302                else:303                    continue304 305            # Enforce unique, monotonically increasing sequence positions per runway306            if runway_id:307                if runway_id not in runway_queue_pos:308                    runway_queue_pos[runway_id] = 1309                seq_pos = runway_queue_pos[runway_id]310                runway_queue_pos[runway_id] += 1311            else:312                seq_pos = action_dict.get("sequence_position", 1)313 314            actions.append(Action(315                aircraft_id=ac_id,316                action_type=action_type,317                runway_id=runway_id,318                sequence_position=seq_pos,319            ))320 321        except Exception:322            pass323 324    return actions325 326 327def run_inference(task_id: int, model_name: str, api_base: str, api_key: str) -> Tuple[float, float, dict]:328    """Run inference on a single task."""329    client = OpenAI(base_url=api_base, api_key=api_key)330    env = ATCEnv(seed=42)331    obs = env.reset(task_id)332    episode_log = {"task_id": task_id, "steps": []}333    all_step_rewards = []334    log_start(task=f"task_{task_id}", env="atc-openenv", model=model_name)335 336    step_num = 0337    total_reward = 0.0338 339    while not obs.episode_done:340        step_num += 1341        state = env.state()342        state_dump = state.model_dump()343        current_step = state_dump.get("timestep", step_num)344        actions = []345        error_msg = None346 347        needs_action = sum(348            1 for ac in state_dump.get("aircraft", [])349            if ac.get("status") in ("INBOUND", "HOLDING")350        )351 352        if needs_action > 0:353            if should_skip_llm(task_id, current_step):354                # Use deterministic controller directly — faster and safer for critical steps355                actions = get_fallback_action(env, task_id=task_id, current_step=current_step)356            else:357                # Use LLM with deterministic fallback358                prompt = build_prompt(obs.model_dump(), state_dump, task_id)359                try:360                    response = client.chat.completions.create(361                        model=model_name,362                        messages=[363                            {364                                "role": "system",365                                "content": (366                                    "You are an expert ATC controller. "367                                    "Respond with valid JSON only. No text outside the JSON block."368                                ),369                            },370                            {"role": "user", "content": prompt},371                        ],372                        temperature=0.0,373                        max_tokens=1500,374                    )375                    response_text = response.choices[0].message.content if response.choices else ""376                    action_dicts = parse_llm_response(response_text)377                    if action_dicts:378                        actions = apply_guardrails(action_dicts, state_dump)379                except Exception as e:380                    error_msg = str(e)381 382                # Fallback if LLM returned nothing or failed383                if not actions:384                    actions = get_fallback_action(env, task_id=task_id, current_step=current_step)385 386        obs, reward, done, info = env.step(action=actions)387        total_reward += reward.value388        all_step_rewards.append(reward.value)389 390        episode_log["steps"].append({"step": step_num, "reward": reward.value})391 392        action_str = (393            json.dumps([a.model_dump() for a in actions]).replace(" ", "")394            if actions else "none"395        )396        log_step(step=step_num, action=action_str, reward=reward.value, done=done, error=error_msg)397 398    grader_inputs = env.get_grader_input()399    episode_log["grader_input"] = grader_inputs400 401    if task_id == 1:402        final_score = grade1(grader_inputs)403    elif task_id == 2:404        final_score = grade2(grader_inputs)405    else:406        final_score = grade3(grader_inputs)407 408    success = final_score >= 0.5409    log_end(success=success, steps=step_num, score=final_score, rewards=all_step_rewards)410 411    return final_score, total_reward, episode_log412 413 414def main():415    API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")416    MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")417    HF_TOKEN = os.getenv("HF_TOKEN")418 419    print(f"[DEBUG] Starting ATC inference with model: {MODEL_NAME}")420    print(f"[DEBUG] API base: {API_BASE_URL}")421 422    total_score = 0.0423    results = []424 425    for task_id in [1, 2, 3]:426        print(f"\n[DEBUG] ===== Running Task {task_id} =====")427        try:428            score, reward, log = run_inference(task_id, MODEL_NAME, API_BASE_URL, HF_TOKEN)429            results.append({"task_id": task_id, "score": score, "reward": reward})430            total_score += score431        except Exception as e:432            print(f"[DEBUG] Error running task {task_id}: {e}")433            results.append({"task_id": task_id, "score": 0.0, "reward": 0.0})434 435    avg_score = total_score / 3.0436 437    print("\n" + "=" * 50)438    print("[DEBUG] FINAL RESULTS")439    print("=" * 50)440    for r in results:441        print(f"[DEBUG] Task {r['task_id']}: score={r['score']:.4f}, reward={r['reward']:.4f}")442    print(f"[DEBUG] Average score: {avg_score:.4f}")443 444 445if __name__ == "__main__":446    main()