updated inference.py

inference.py

@@ -1,33 +1,31 @@
 """
-inference.py — Space Manufacturing RL submission entry point.
-
-Default policy: OpenAI (falls back to heuristic if the client cannot be built).
+inference.py — Satellite constellation RL submission entry point.
 
 Environment variables:
-  OPENAI_API_KEY    — API key (required for OpenAI policy); alias: API_KEY
-  OPENAI_BASE_URL   — Custom base URL (optional, e.g. Azure / proxy); alias: API_BASE_URL
-  MODEL_NAME        — Model to use (default: gpt-4o-mini)
-  BASELINE_POLICY   — Force policy: "openai" (default) or "heuristic"
-  TEMPERATURE       — Sampling temperature (default: 0.0)
-  MAX_TOKENS        — Max tokens per response (default: 300)
-  REQUEST_DELAY     — Seconds to sleep between steps (default: 0.0)
-  REQUEST_TIMEOUT   — HTTP timeout in seconds (default: 30.0)
-  DEBUG             — Print raw model responses when "true"
+  API_BASE_URL     — OpenAI-compatible endpoint base URL (required)
+  API_KEY          — API key (required)
+  MODEL_NAME       — Model to use (required)
+  BASELINE_POLICY  — Force policy: "openai" (default) or "heuristic"
+  TEMPERATURE      — Sampling temperature (default: 0.0)
+  MAX_TOKENS       — Max tokens per response (default: 300)
+  REQUEST_DELAY    — Seconds to sleep between steps (default: 0.0)
+  REQUEST_TIMEOUT  — HTTP timeout in seconds (default: 30.0)
+  STEP_TIMEOUT     — Per-step inference wall-clock timeout in seconds (default: 45.0)
+  TASK_TIMEOUT     — Per-task wall-clock timeout in seconds, 0 = no limit (default: 0.0)
+  DEBUG            — Print raw model responses when "true"
 
 Usage:
-  # OpenAI (default):
-  OPENAI_API_KEY=sk-... python inference.py
-
-  # Custom base URL / compatible provider:
-  OPENAI_API_KEY=hf_... OPENAI_BASE_URL=https://... MODEL_NAME=mistralai/Mixtral-8x7B-Instruct-v0.1 python inference.py
+  API_BASE_URL=https://... API_KEY=hf_... MODEL_NAME=mistralai/... python inference.py
 
-  # Force heuristic baseline:
-  BASELINE_POLICY=heuristic python inference.py
+  # Per-step and per-task timeouts:
+  STEP_TIMEOUT=20 TASK_TIMEOUT=300 python inference.py
 """
 from __future__ import annotations
 
 import asyncio
+import importlib.util
 import json
+import math
 import os
 import re
 import sys
@@ -35,31 +33,36 @@ import textwrap
 import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 from urllib.parse import urlparse
 
-try:
-    from openai import OpenAI
-except ImportError:  # pragma: no cover
-    OpenAI = None  # type: ignore[assignment,misc]
+from openai import OpenAI
 
 try:
     from dotenv import load_dotenv
 except Exception:  # pragma: no cover
     load_dotenv = None  # type: ignore[assignment]
 
-# ── package imports ────────────────────────────────────────────────────────────
-# inference.py is a script inside SpaceFactory/. Insert the parent directory so
-# the whole folder is importable as the 'SpaceFactory' package, which keeps all
-# relative imports inside the package working correctly.
-_pkg_parent = str(Path(__file__).resolve().parent.parent)
-if _pkg_parent not in sys.path:
-    sys.path.insert(0, _pkg_parent)
-
-from SpaceFactory.env import ManufacturingTaskEnv
-from SpaceFactory.graders import ManufacturingTaskGrader
-from SpaceFactory.models import ManufacturingAction, ManufacturingObservation
-from SpaceFactory.tasks import EasyTask, HardTask, MediumTask
+try:
+    from satellite import EasyTask, HardTask, MediumTask, SatelliteAction, SatelliteTaskEnv, TaskGrader
+except ImportError:  # pragma: no cover
+    package_root = Path(__file__).resolve().parent
+    spec = importlib.util.spec_from_file_location(
+        "satellite",
+        package_root / "__init__.py",
+        submodule_search_locations=[str(package_root)],
+    )
+    if spec is None or spec.loader is None:
+        raise
+    satellite = importlib.util.module_from_spec(spec)
+    sys.modules["satellite"] = satellite
+    spec.loader.exec_module(satellite)
+    EasyTask = satellite.EasyTask
+    HardTask = satellite.HardTask
+    MediumTask = satellite.MediumTask
+    SatelliteAction = satellite.SatelliteAction
+    SatelliteTaskEnv = satellite.SatelliteTaskEnv
+    TaskGrader = satellite.TaskGrader
 
 if load_dotenv is not None:
     load_dotenv()
@@ -100,198 +103,346 @@ def read_int_env(name: str, default: int) -> int:
 
 
 # ── configuration ──────────────────────────────────────────────────────────────
-API_KEY         = os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY") or os.getenv("HF_TOKEN")
-API_BASE_URL    = os.getenv("OPENAI_BASE_URL") or os.getenv("API_BASE_URL")
-MODEL_NAME      = os.getenv("MODEL_NAME", "gpt-4o-mini")
+API_BASE_URL    = os.environ["API_BASE_URL"]
+API_KEY         = os.environ["API_KEY"]
+MODEL_NAME      = os.getenv("MODEL_NAME","meta-llama/Llama-3.1-8B-Instruct:novita")
 BASELINE_POLICY = os.getenv("BASELINE_POLICY", "openai").lower()
 TEMPERATURE     = read_float_env("TEMPERATURE", 0.0)
 MAX_TOKENS      = read_int_env("MAX_TOKENS", 300)
 REQUEST_DELAY   = read_float_env("REQUEST_DELAY", 0.0)
 REQUEST_TIMEOUT = read_float_env("REQUEST_TIMEOUT", 30.0)
+STEP_TIMEOUT    = read_float_env("STEP_TIMEOUT", 45.0)  # per-step wall-clock limit
+TASK_TIMEOUT    = read_float_env("TASK_TIMEOUT", 0.0)   # per-task limit; 0 = no limit
 DEBUG           = os.getenv("DEBUG", "false").lower() == "true"
 
-FALLBACK_ACTION = "recharge"
+FALLBACK_ACTION = "idle"
 TASK_ORDER      = ["easy", "medium", "hard"]
 TASK_TYPES      = {"easy": EasyTask, "medium": MediumTask, "hard": HardTask}
 
-VALID_ACTIONS   = {"produce", "assemble", "deliver", "recharge"}
-ACTION_PATTERN  = re.compile(r"(produce|assemble|deliver|recharge)", re.IGNORECASE)
+ACTION_PATTERN  = re.compile(r"(capture|downlink|maintain|idle)", re.IGNORECASE)
+ACTION_PREFIX_RE = re.compile(r"^(action|next action)\s*[:\-]\s*", re.IGNORECASE)
 
 # ── system prompt ──────────────────────────────────────────────────────────────
-SYSTEM_PROMPT = textwrap.dedent("""
-    You are controlling orbital manufacturing platforms.
-    Each step, output ONLY a JSON object mapping platform IDs (as strings) to one of:
-      "produce", "assemble", "deliver", "recharge"
+SYSTEM_PROMPT = textwrap.dedent(
+    """
+    You are managing a real-world satellite constellation.
+    Reply with exactly one JSON object mapping satellite ids to actions.
+
+    Valid actions:
+    - capture
+    - downlink
+    - maintain
+    - idle
 
     Decision guidance:
-    - recharge immediately if energy < 15
-    - deliver when product_stock > 0 and a delivery window is open
-    - assemble when component_stock >= 10 and product_stock < 5
-    - produce when material_stock >= 15 and component_stock < 30
-    - recharge when energy < 40 and no urgent action is available
-    - avoid invalid actions (e.g. assemble with no components)
-    - keep all platforms energy-healthy across the full episode
+    - prefer capture only when a visible image task exists and the satellite has battery/storage margin
+    - prefer downlink when a visible ground station task exists, especially if storage is high
+    - use maintain to recover low-battery satellites before they become risky
+    - avoid invalid, repeated, or wasteful actions
+    - keep the fleet healthy across the full episode, not just the current step
 
-    Output format (no explanation, no markdown):
-    {"0": "produce", "1": "assemble", "2": "deliver"}
-""").strip()
+    Output format:
+    {"0": "capture", "1": "idle"}
 
+    Do not include explanations or any extra text outside the JSON object.
+    """
+).strip()
 
-# ── result dataclass ──────────────────────────────────────────────────────────
+
+# ── result dataclass ───────────────────────────────────────────────────────────
 @dataclass
 class TaskRunResult:
     task_name: str
-    score: float
-    total_reward: float
+    grade: float                          # 0.0 – 1.0  (primary output)
+    grade_components: Dict[str, float]    # per-criterion, each 0.0 – 1.0
     steps: int
     done: bool
-    metrics: Dict[str, float]
+    metrics: Dict[str, float]             # raw env metrics (informational)
 
 
 # ── observation helpers ────────────────────────────────────────────────────────
 
-def observation_to_dict(obs: ManufacturingObservation) -> Dict[str, Any]:
+def build_history_lines(history: List[str]) -> str:
+    if not history:
+        return "None"
+    return "\n".join(history[-6:])
+
+
+def satellite_geo(position: Tuple[float, float, float]) -> Tuple[float, float, float]:
+    x, y, z = position
+    radius = math.sqrt((x * x) + (y * y) + (z * z))
+    if radius <= 0:
+        return 0.0, 0.0, 0.0
+    lat = math.degrees(math.asin(z / radius))
+    lon = math.degrees(math.atan2(y, x))
+    altitude = max(0.0, radius - 6371.0)
+    return float(lat), float(lon), float(altitude)
+
+
+def visibility_radius_rad(altitude_km: float) -> float:
+    earth_radius_km = 6371.0
+    alt = max(0.0, altitude_km)
+    horizon = math.acos(min(1.0, earth_radius_km / (earth_radius_km + alt)))
+    return max(math.radians(35.0), min(math.radians(120.0), horizon + math.radians(50.0)))
+
+
+def great_circle_distance_rad(lat1: float, lon1: float, lat2: float, lon2: float) -> float:
+    lat1_rad = math.radians(lat1)
+    lon1_rad = math.radians(lon1)
+    lat2_rad = math.radians(lat2)
+    lon2_rad = math.radians(lon2)
+    d_lat = lat2_rad - lat1_rad
+    d_lon = lon2_rad - lon1_rad
+    a = (
+        math.sin(d_lat / 2.0) ** 2
+        + math.cos(lat1_rad) * math.cos(lat2_rad) * math.sin(d_lon / 2.0) ** 2
+    )
+    return 2.0 * math.asin(min(1.0, math.sqrt(a)))
+
+
+def normalize_pending_tasks(tasks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    normalized = []
+    for task in tasks:
+        normalized.append({key: value for key, value in task.items()})
+    return normalized
+
+
+def extract_capture_regions(observation: Dict[str, Any]) -> Dict[str, Tuple[float, float]]:
+    regions = observation.get("capture_regions")
+    if isinstance(regions, dict) and regions:
+        return {
+            str(name): (float(coords[0]), float(coords[1]))
+            for name, coords in regions.items()
+            if isinstance(coords, (list, tuple)) and len(coords) == 2
+        }
     return {
-        "platforms": [
+        "region1": (18.5, 73.9),
+        "region2": (34.0, -117.0),
+        "region3": (-22.8, -43.2),
+    }
+
+
+def find_visible_capture_task(
+    sat: Dict[str, Any],
+    observation: Dict[str, Any],
+) -> Optional[Dict[str, Any]]:
+    capture_regions = extract_capture_regions(observation)
+    sat_lat, sat_lon, sat_alt = satellite_geo(tuple(sat["position"]))
+    max_distance = visibility_radius_rad(sat_alt)
+    candidates: List[Tuple[float, str, Dict[str, Any]]] = []
+    weather = observation.get("weather_conditions", {})
+
+    for task in normalize_pending_tasks(observation.get("pending_tasks", [])):
+        if task.get("type") != "image_capture":
+            continue
+        region = str(task.get("region", ""))
+        if region not in capture_regions:
+            continue
+        reg_lat, reg_lon = capture_regions[region]
+        distance = great_circle_distance_rad(sat_lat, sat_lon, reg_lat, reg_lon)
+        if distance > max_distance:
+            continue
+        priority = float(task.get("priority", 1))
+        cloud = float(weather.get(region, 0.5))
+        score = (priority * 3.0) + ((1.0 - cloud) * 2.0) - distance
+        candidates.append((score, str(task.get("id", "")), task))
+
+    if not candidates:
+        return None
+    candidates.sort(key=lambda item: (-item[0], item[1]))
+    return candidates[0][2]
+
+
+def find_visible_downlink_task(
+    sat: Dict[str, Any],
+    observation: Dict[str, Any],
+) -> Optional[Dict[str, Any]]:
+    stations = observation.get("ground_stations", [])
+    sat_lat, sat_lon, sat_alt = satellite_geo(tuple(sat["position"]))
+    max_distance = visibility_radius_rad(sat_alt)
+    candidates: List[Tuple[float, str, Dict[str, Any]]] = []
+
+    for task in normalize_pending_tasks(observation.get("pending_tasks", [])):
+        if task.get("type") != "data_downlink":
+            continue
+        station_id = int(task.get("station", 0))
+        if station_id < 0 or station_id >= len(stations):
+            continue
+        gs_lat, gs_lon = stations[station_id]
+        distance = great_circle_distance_rad(sat_lat, sat_lon, float(gs_lat), float(gs_lon))
+        if distance > max_distance:
+            continue
+        priority = float(task.get("priority", 1))
+        units_remaining = float(task.get("units_remaining", 20.0))
+        completion_bias = 0.75 if float(sat["storage"]) >= units_remaining else 0.0
+        score = (priority * 3.0) + completion_bias - distance
+        candidates.append((score, str(task.get("id", "")), task))
+
+    if not candidates:
+        return None
+    candidates.sort(key=lambda item: (-item[0], item[1]))
+    return candidates[0][2]
+
+
+def observation_to_dict(observation: Any) -> Dict[str, Any]:
+    return {
+        "satellites": [
             {
-                "id": p.id,
-                "energy": p.energy,
-                "material_stock": p.material_stock,
-                "component_stock": p.component_stock,
-                "product_stock": p.product_stock,
-                "last_action": p.last_action,
+                "id": sat.id,
+                "position": sat.position,
+                "battery": sat.battery,
+                "storage": sat.storage,
+                "last_action": sat.last_action,
             }
-            for p in obs.platforms
-        ],
-        "time_step": obs.time_step,
-        "delivery_windows": [
-            {"order_id": w.order_id, "product_type": w.product_type, "deadline": w.deadline}
-            for w in obs.delivery_windows
+            for sat in observation.satellites
         ],
-        "solar_conditions": obs.solar_conditions,
-        "pending_orders": [
-            {"order_id": o.order_id, "product_type": o.product_type,
-             "requires_assembly": o.requires_assembly}
-            for o in obs.pending_orders
-        ],
-        "total_reward": obs.total_reward,
-        "done": obs.done,
-        "reward": obs.reward,
-        "metadata": obs.metadata,
+        "time_step": observation.time_step,
+        "ground_stations": observation.ground_stations,
+        "weather_conditions": observation.weather_conditions,
+        "pending_tasks": observation.pending_tasks,
+        "total_reward": observation.total_reward,
+        "done": observation.done,
+        "reward": observation.reward,
+        "metadata": observation.metadata,
     }
 
 
-def build_idle_actions(obs_dict: Dict[str, Any]) -> Dict[int, str]:
-    return {int(p["id"]): FALLBACK_ACTION for p in obs_dict.get("platforms", [])}
+def build_idle_actions(observation: Dict[str, Any]) -> Dict[int, str]:
+    actions: Dict[int, str] = {}
+    for sat in observation.get("satellites", []):
+        try:
+            actions[int(sat["id"])] = FALLBACK_ACTION
+        except (KeyError, TypeError, ValueError):
+            continue
+    return actions
 
 
 # ── heuristic policy ───────────────────────────────────────────────────────────
 
-def heuristic_action(obs_dict: Dict[str, Any]) -> Dict[int, str]:
+def heuristic_action(observation: Dict[str, Any]) -> Dict[int, str]:
     actions: Dict[int, str] = {}
-    has_open_window = len(obs_dict.get("delivery_windows", [])) > 0
-
-    for p in obs_dict.get("platforms", []):
-        pid     = int(p["id"])
-        energy  = float(p["energy"])
-        mat     = float(p["material_stock"])
-        comp    = float(p["component_stock"])
-        prod    = int(p["product_stock"])
-
-        if energy < 15.0:
-            action = "recharge"
-        elif prod > 0 and has_open_window:
-            action = "deliver"
-        elif comp >= 10.0 and prod < 5:
-            action = "assemble"
-        elif mat >= 15.0 and comp < 30.0:
-            action = "produce"
-        elif energy < 40.0:
-            action = "recharge"
-        elif mat >= 15.0:
-            action = "produce"
-        else:
-            action = "recharge"
+    pending_tasks = observation.get("pending_tasks", [])
+    has_capture_task = any(task.get("type") == "image_capture" for task in pending_tasks)
+    has_downlink_task = any(task.get("type") == "data_downlink" for task in pending_tasks)
+
+    for sat in observation.get("satellites", []):
+        sat_id = int(sat["id"])
+        battery = float(sat["battery"])
+        storage = float(sat["storage"])
+        visible_capture = find_visible_capture_task(sat, observation)
+        visible_downlink = find_visible_downlink_task(sat, observation)
+
+        if battery <= 12:
+            actions[sat_id] = "maintain"
+            continue
+        if battery < 28 and not visible_downlink:
+            actions[sat_id] = "maintain"
+            continue
+        if storage >= 85 and visible_downlink:
+            actions[sat_id] = "downlink"
+            continue
+        if visible_capture and battery >= 25 and storage <= 80:
+            actions[sat_id] = "capture"
+            continue
+        if visible_downlink and storage > 0:
+            actions[sat_id] = "downlink"
+            continue
+        if battery < 45 and not has_capture_task:
+            actions[sat_id] = "maintain"
+            continue
+        if battery < 35 and storage <= 5:
+            actions[sat_id] = "maintain"
+            continue
+        if has_downlink_task and storage >= 50:
+            actions[sat_id] = "idle"
+            continue
+        if has_capture_task and battery >= 30 and storage < 70:
+            actions[sat_id] = "idle"
+            continue
+        actions[sat_id] = "idle"
 
-        actions[pid] = action
     return actions
 
 
-def safe_heuristic_action(obs_dict: Dict[str, Any], reason: str) -> Dict[int, str]:
+def safe_heuristic_action(observation: Dict[str, Any], reason: str) -> Dict[int, str]:
     try:
-        return heuristic_action(obs_dict)
+        return heuristic_action(observation)
     except Exception as exc:  # noqa: BLE001
         warn_once(
             f"heuristic:{reason}",
-            f"Heuristic fallback failed after {reason}: {exc}. Returning all-{FALLBACK_ACTION}.",
+            f"Heuristic fallback failed after {reason}: {exc}. Returning all-idle actions.",
         )
-        return build_idle_actions(obs_dict)
+        return build_idle_actions(observation)
 
 
 # ── prompt formatting ──────────────────────────────────────────────────────────
 
-def build_history_lines(history: List[str]) -> str:
-    return "\n".join(history[-6:]) if history else "None"
-
-
-def format_observation(task_name: str, obs_dict: Dict[str, Any]) -> str:
-    platforms_lines = []
-    for p in obs_dict.get("platforms", []):
-        platforms_lines.append(
-            f"  [{p['id']}] energy={p['energy']:.1f}  mat={p['material_stock']:.1f}"
-            f"  comp={p['component_stock']:.1f}  prod={p['product_stock']}"
-            f"  last={p['last_action']}"
-        )
-
-    windows_lines = []
-    for w in obs_dict.get("delivery_windows", []):
-        step_now = obs_dict.get("time_step", 0)
-        urgency  = w["deadline"] - step_now
-        windows_lines.append(
-            f"  order={w['order_id']} type={w['product_type']}"
-            f" deadline={w['deadline']} ({urgency} steps left)"
+def format_observation(task_name: str, observation: Dict[str, Any]) -> str:
+    satellites_info = []
+    for sat in observation.get("satellites", []):
+        sat_lat, sat_lon, sat_alt = satellite_geo(tuple(sat["position"]))
+        capture_task = find_visible_capture_task(sat, observation)
+        downlink_task = find_visible_downlink_task(sat, observation)
+        satellites_info.append(
+            f"  Satellite {sat['id']}: battery={sat['battery']:.1f}, "
+            f"storage={sat['storage']:.1f}, last={sat['last_action']}, "
+            f"lat={sat_lat:.1f}, lon={sat_lon:.1f}, alt={sat_alt:.1f}km, "
+            f"capture_visible={capture_task is not None}, "
+            f"downlink_visible={downlink_task is not None}"
         )
 
-    solar_str = ", ".join(
-        f"{z}={round(v * 100)}%"
-        for z, v in obs_dict.get("solar_conditions", {}).items()
+    tasks_info = []
+    for task in observation.get("pending_tasks", [])[:12]:
+        descriptor = task["type"]
+        if task["type"] == "image_capture":
+            descriptor += f" region={task.get('region')}"
+        if task["type"] == "data_downlink":
+            descriptor += f" station={task.get('station')}"
+            descriptor += f" units={task.get('units_remaining', 0)}"
+        tasks_info.append(f"  - {descriptor} priority={task.get('priority', 1)}")
+
+    weather_info = ", ".join(
+        f"{region}={cover:.0%}"
+        for region, cover in observation.get("weather_conditions", {}).items()
     )
 
-    return textwrap.dedent(f"""
+    return textwrap.dedent(
+        f"""
         Task: {task_name}
-        Time Step: {obs_dict.get('time_step', 0)}
-        Total Reward: {obs_dict.get('total_reward', 0.0):.2f}
-        Solar: {solar_str or 'n/a'}
-
-        Platforms:
-        {chr(10).join(platforms_lines) or '  None'}
+        Time Step: {observation.get('time_step', 0)}
+        Total Reward: {observation.get('total_reward', 0.0):.2f}
+        Weather: {weather_info or 'n/a'}
 
-        Open Delivery Windows:
-        {chr(10).join(windows_lines) if windows_lines else '  None'}
+        Satellites:
+        {chr(10).join(satellites_info) or '  None'}
 
-        Pending orders: {len(obs_dict.get('pending_orders', []))}
-    """).strip()
+        Pending Tasks:
+        {chr(10).join(tasks_info) or '  - None'}
+        """
+    ).strip()
 
 
 def build_user_prompt(
     task_name: str,
     step: int,
-    obs_dict: Dict[str, Any],
+    observation: Dict[str, Any],
     history: List[str],
     total_reward: float,
 ) -> str:
-    return textwrap.dedent(f"""
+    return textwrap.dedent(
+        f"""
         Step: {step}
         Aggregate reward so far: {total_reward:+.2f}
 
         Current state:
-        {format_observation(task_name, obs_dict)}
+        {format_observation(task_name, observation)}
 
         Previous steps:
         {build_history_lines(history)}
 
         Reply with exactly one JSON object.
-    """).strip()
+        """
+    ).strip()
 
 
 # ── model response parsing ────────────────────────────────────────────────────
@@ -309,43 +460,46 @@ def extract_response_text(completion: Any) -> str:
     if isinstance(content, list):
         parts: List[str] = []
         for item in content:
-            text = item.get("text") if isinstance(item, dict) else getattr(item, "text", None)
+            if isinstance(item, dict):
+                text = item.get("text")
+            else:
+                text = getattr(item, "text", None)
             if text:
                 parts.append(str(text))
         return "\n".join(parts)
     return str(content or "")
 
 
-def parse_model_action(
-    response_text: str, obs_dict: Dict[str, Any]
-) -> Dict[int, str]:
+def parse_model_action(response_text: str, observation: Dict[str, Any]) -> Dict[int, str]:
     if not response_text:
-        return safe_heuristic_action(obs_dict, "empty model response")
+        return safe_heuristic_action(observation, "empty model response")
+
+    cleaned = ACTION_PREFIX_RE.sub("", response_text.strip())
 
     try:
-        json_match = re.search(r"\{.*\}", response_text.strip(), re.DOTALL)
+        json_match = re.search(r"\{.*\}", cleaned, re.DOTALL)
         if json_match:
             parsed = json.loads(json_match.group(0))
-            valid_ids = {int(p["id"]) for p in obs_dict.get("platforms", [])}
             actions: Dict[int, str] = {}
+            valid_ids = {int(sat["id"]) for sat in observation.get("satellites", [])}
             for key, value in parsed.items():
-                pid    = int(key)
-                if pid not in valid_ids:
+                sat_id = int(key)
+                if sat_id not in valid_ids:
                     continue
                 action = str(value).strip().lower()
-                if action not in VALID_ACTIONS:
+                if not ACTION_PATTERN.fullmatch(action):
                     action = FALLBACK_ACTION
-                actions[pid] = action
+                actions[sat_id] = action
             if actions:
-                fallback = heuristic_action(obs_dict)
-                for pid in valid_ids:
-                    actions.setdefault(pid, fallback.get(pid, FALLBACK_ACTION))
+                fallback = heuristic_action(observation)
+                for sat_id in valid_ids:
+                    actions.setdefault(sat_id, fallback.get(sat_id, FALLBACK_ACTION))
                 return actions
     except (json.JSONDecodeError, TypeError, ValueError):
         pass
 
-    warn_once("parse:model-response", "Model response was not valid JSON; using heuristic.")
-    return safe_heuristic_action(obs_dict, "invalid model response")
+    warn_once("parse:model-response", "Model response was not valid JSON; using heuristic actions.")
+    return safe_heuristic_action(observation, "invalid model response")
 
 
 # ── client construction ────────────────────────────────────────────────────────
@@ -354,7 +508,7 @@ def validate_api_base_url(base_url: Optional[str]) -> Optional[str]:
     if not base_url:
         return None
     cleaned = base_url.strip().rstrip("/")
-    parsed  = urlparse(cleaned)
+    parsed = urlparse(cleaned)
     if parsed.scheme not in {"http", "https"} or not parsed.netloc:
         warn_once(
             "config:api-base-url",
@@ -364,48 +518,30 @@ def validate_api_base_url(base_url: Optional[str]) -> Optional[str]:
     return cleaned
 
 
-def build_client() -> Optional[Any]:
-    if BASELINE_POLICY == "heuristic":
-        return None
-
-    if OpenAI is None:
-        warn_once("client:import", "openai package not installed; falling back to heuristic policy.")
-        return None
-
-    if not API_KEY:
-        warn_once(
-            "config:missing",
-            "OPENAI_API_KEY is not set. Falling back to heuristic policy.",
-        )
-        return None
-
-    validated_base = validate_api_base_url(API_BASE_URL)  # None = use OpenAI default endpoint
-
-    try:
-        kwargs: Dict[str, Any] = {"api_key": API_KEY, "timeout": REQUEST_TIMEOUT}
-        if validated_base:
-            kwargs["base_url"] = validated_base
-        return OpenAI(**kwargs)
-    except Exception as exc:  # noqa: BLE001
-        warn_once("client:init", f"Failed to build OpenAI client: {exc}. Using heuristic.")
-        return None
+def build_client() -> OpenAI:
+    return OpenAI(
+        base_url=API_BASE_URL,
+        api_key=API_KEY,
+        timeout=REQUEST_TIMEOUT,
+    )
 
 
 # ── action chooser ─────────────────────────────────────────────────────────────
 
-def choose_actions(
-    client: Optional[Any],
+async def choose_actions(
+    client: Optional[OpenAI],
     task_name: str,
     step: int,
-    obs_dict: Dict[str, Any],
+    observation: Dict[str, Any],
     history: List[str],
     total_reward: float,
 ) -> Dict[int, str]:
     if client is None:
-        return safe_heuristic_action(obs_dict, "heuristic mode")
+        raise RuntimeError("OpenAI client not initialized")
 
-    user_prompt = build_user_prompt(task_name, step, obs_dict, history, total_reward)
-    try:
+    user_prompt = build_user_prompt(task_name, step, observation, history, total_reward)
+
+    def _call() -> str:
         completion = client.chat.completions.create(
             model=MODEL_NAME,
             messages=[
@@ -415,79 +551,95 @@ def choose_actions(
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
         )
-        response_text = extract_response_text(completion)
+        return extract_response_text(completion)
+
+    try:
+        loop = asyncio.get_event_loop()
+        response_text = await asyncio.wait_for(
+            loop.run_in_executor(None, _call),
+            timeout=STEP_TIMEOUT,
+        )
+    except asyncio.TimeoutError:
+        warn_once(
+            f"timeout:{task_name}",
+            f"[{task_name}] Step {step} timed out after {STEP_TIMEOUT}s. Using heuristic actions.",
+        )
+        return safe_heuristic_action(observation, f"step timeout on {task_name} step {step}")
     except Exception as exc:  # noqa: BLE001
         warn_once(
             f"model:{task_name}",
-            f"[{task_name}] Model request failed at step {step}: {exc}. Using heuristic.",
+            f"[{task_name}] Model request failed at step {step}: {exc}. Using heuristic actions.",
         )
-        return safe_heuristic_action(obs_dict, f"model failure on {task_name} step {step}")
+        return safe_heuristic_action(observation, f"model request failure on {task_name} step {step}")
 
     if DEBUG:
-        print(f"[DEBUG] [{task_name}] step={step} model_response={response_text[:300]!r}")
+        print(f"[{task_name}] model response: {response_text[:300]}")
 
     try:
-        return parse_model_action(response_text, obs_dict)
+        return parse_model_action(response_text, observation)
     except Exception as exc:  # noqa: BLE001
         warn_once(
             f"parse:{task_name}",
-            f"[{task_name}] Parse failed at step {step}: {exc}. Using heuristic.",
+            f"[{task_name}] Failed to parse model response at step {step}: {exc}. "
+            "Using heuristic actions.",
         )
-        return safe_heuristic_action(obs_dict, f"parse failure on {task_name} step {step}")
+        return safe_heuristic_action(observation, f"parse failure on {task_name} step {step}")
 
 
 # ── episode runner ─────────────────────────────────────────────────────────────
 
-async def run_task(task_name: str, client: Optional[Any]) -> TaskRunResult:
-    env    = ManufacturingTaskEnv(task_name=task_name)
-    grader = ManufacturingTaskGrader(task_name=task_name)
+async def run_task(task_name: str, client: Optional[OpenAI]) -> TaskRunResult:
+    env = SatelliteTaskEnv(task_name=task_name)
+    task = TASK_TYPES[task_name]()
+    grader = TaskGrader(task)
     history: List[str] = []
 
-    obs        = env.reset()
-    state      = env.state()
+    observation = env.reset()
+    state = env.state()
     step_limit = state.max_steps
-
-    print(f"[START] task={task_name} max_steps={step_limit}", flush=True)
+    print(f"[START] task={task_name}", flush=True)
 
     for step in range(1, step_limit + 1):
-        obs_dict = observation_to_dict(obs)
-        actions  = choose_actions(client, task_name, step, obs_dict, history, obs.total_reward)
-
-        # Wrap dict back into ManufacturingAction
-        action_obj = ManufacturingAction(platform_actions=actions)
-        obs, reward, done, info = env.step(action_obj)
+        obs_dict = observation_to_dict(observation)
+        actions = await choose_actions(
+            client,
+            task_name,
+            step,
+            obs_dict,
+            history,
+            observation.total_reward,
+        )
 
+        observation, reward, done, info = env.step(
+            SatelliteAction(satellite_actions=actions)
+        )
         reward_value = float(reward.value)
         history.append(f"step {step}: {actions} -> reward {reward_value:+.2f}")
 
         print(
-            f"[STEP]  task={task_name} step={step}/{step_limit}"
-            f" reward={reward_value:.4f} total={obs.total_reward:.4f}"
-            f" done={done}",
+            f"[STEP] step={step} reward={reward_value:.4f} total={observation.total_reward:.4f}",
             flush=True,
         )
 
-        if REQUEST_DELAY > 0 and not done:
+        if REQUEST_DELAY > 0 and step < step_limit and not done:
             time.sleep(REQUEST_DELAY)
 
         if done:
             break
 
     final_state = env.state()
-    metrics     = {k: float(v) for k, v in final_state.metrics.items()}
-    score       = grader.grade(metrics, final_state.step_count, final_state.platforms)
-
+    metrics = {key: float(value) for key, value in final_state.metrics.items()}
+    grade = grader.grade_episode(env)
+    grade_components = grader.grade_components(env)
     print(
-        f"[END]   task={task_name} score={score:.4f}"
-        f" steps={final_state.step_count} total_reward={final_state.total_reward:.4f}"
-        f" done={final_state.done}",
+        f"[END] task={task_name} grade={grade:.4f} "
+        f"steps={final_state.step_count} done={final_state.done}",
         flush=True,
     )
-
     return TaskRunResult(
         task_name=task_name,
-        score=score,
-        total_reward=final_state.total_reward,
+        grade=grade,
+        grade_components=grade_components,
         steps=final_state.step_count,
         done=final_state.done,
         metrics=metrics,
@@ -497,29 +649,64 @@ async def run_task(task_name: str, client: Optional[Any]) -> TaskRunResult:
 # ── summary printer ────────────────────────────────────────────────────────────
 
 def print_summary(results: List[TaskRunResult]) -> None:
-    aggregate = sum(r.score for r in results) / len(results)
-    print("\nInference Summary")
-    print("=" * 60)
+    aggregate = sum(r.grade for r in results) / len(results)
+    all_keys: list = []
     for r in results:
-        print(
-            f"{r.task_name:<8} score={r.score:.4f}"
-            f"  reward={r.total_reward:.2f}"
-            f"  steps={r.steps}"
-            f"  done={r.done}"
-        )
-    print("-" * 60)
-    print(f"aggregate_score={aggregate:.4f}")
-    print("=" * 60)
+        for k in r.grade_components:
+            if k not in all_keys:
+                all_keys.append(k)
+
+    label_map = getattr(
+        __import__("graders", fromlist=["TaskGrader"]).TaskGrader, "CRITERION_LABELS", {}
+    )
+
+    col_w = 10
+    header_parts = [f"{'task':<8}", f"{'grade':>7}"]
+    for k in all_keys:
+        label = label_map.get(k, k)[:col_w]
+        header_parts.append(f"{label:>{col_w}}")
+    header_parts.append(f"{'steps':>6}")
+
+    sep_width = 8 + 7 + col_w * len(all_keys) + 6 + len(all_keys) * 2 + 10
+    print("\nInference Grade Summary  (all values 0.0 – 1.0)")
+    print("=" * sep_width)
+    print("  ".join(header_parts))
+    print("-" * sep_width)
+    for r in results:
+        row = [f"{r.task_name:<8}", f"{r.grade:>7.4f}"]
+        for k in all_keys:
+            v = r.grade_components.get(k, float("nan"))
+            row.append(f"{v:>{col_w}.4f}")
+        row.append(f"{r.steps:>6}")
+        print("  ".join(row))
+    print("-" * sep_width)
+    print(f"  {'aggregate':<8}  {aggregate:>7.4f}")
+    print("=" * sep_width)
 
 
 # ── async main ─────────────────────────────────────────────────────────────────
 
 async def async_main() -> None:
-    client  = build_client()
+    client = build_client()
     results = []
     for task_name in TASK_ORDER:
-        results.append(await run_task(task_name, client))
-    print_summary(results)
+        if TASK_TIMEOUT > 0:
+            try:
+                result = await asyncio.wait_for(
+                    run_task(task_name, client), timeout=TASK_TIMEOUT
+                )
+            except asyncio.TimeoutError:
+                print(
+                    f"[TIMEOUT] task={task_name} exceeded {TASK_TIMEOUT}s; skipping.",
+                    file=sys.stderr,
+                    flush=True,
+                )
+                continue
+        else:
+            result = await run_task(task_name, client)
+        results.append(result)
+    if results:
+        print_summary(results)
 
 
 def main() -> None: