Reverted and updated Inference.py

inference.py

@@ -1,5 +1,7 @@
 """
-inference.py — Satellite constellation RL submission entry point.
+inference.py — Space Manufacturing RL submission entry point.
+
+Default policy: OpenAI (falls back to heuristic if the client cannot be built).
 
 Environment variables:
   API_BASE_URL     — OpenAI-compatible endpoint base URL (required)
@@ -17,52 +19,47 @@ Environment variables:
 Usage:
   API_BASE_URL=https://... API_KEY=hf_... MODEL_NAME=mistralai/... python inference.py
 
-  # Per-step and per-task timeouts:
+  # Force heuristic baseline:
+  BASELINE_POLICY=heuristic python inference.py
+
+  # 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
 import textwrap
-import time
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 from urllib.parse import urlparse
 
-from openai import OpenAI
+try:
+    from openai import OpenAI
+except ImportError:  # pragma: no cover
+    OpenAI = None  # type: ignore[assignment,misc]
 
 try:
     from dotenv import load_dotenv
 except Exception:  # pragma: no cover
     load_dotenv = None  # type: ignore[assignment]
 
-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
+# ── 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
 
 if load_dotenv is not None:
     load_dotenv()
@@ -103,9 +100,9 @@ def read_int_env(name: str, default: int) -> int:
 
 
 # ── configuration ──────────────────────────────────────────────────────────────
-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")
+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")
 BASELINE_POLICY = os.getenv("BASELINE_POLICY", "openai").lower()
 TEMPERATURE     = read_float_env("TEMPERATURE", 0.0)
 MAX_TOKENS      = read_int_env("MAX_TOKENS", 300)
@@ -115,334 +112,195 @@ STEP_TIMEOUT    = read_float_env("STEP_TIMEOUT", 45.0)  # per-step wall-clock li
 TASK_TIMEOUT    = read_float_env("TASK_TIMEOUT", 0.0)   # per-task limit; 0 = no limit
 DEBUG           = os.getenv("DEBUG", "false").lower() == "true"
 
-FALLBACK_ACTION = "idle"
+FALLBACK_ACTION = "recharge"
 TASK_ORDER      = ["easy", "medium", "hard"]
 TASK_TYPES      = {"easy": EasyTask, "medium": MediumTask, "hard": HardTask}
 
-ACTION_PATTERN  = re.compile(r"(capture|downlink|maintain|idle)", re.IGNORECASE)
-ACTION_PREFIX_RE = re.compile(r"^(action|next action)\s*[:\-]\s*", re.IGNORECASE)
+VALID_ACTIONS   = {"produce", "assemble", "deliver", "recharge"}
+ACTION_PATTERN  = re.compile(r"(produce|assemble|deliver|recharge)", re.IGNORECASE)
 
-# ── system prompt ──────────────────────────────────────────────────────────────
-SYSTEM_PROMPT = textwrap.dedent(
-    """
-    You are managing a real-world satellite constellation.
-    Reply with exactly one JSON object mapping satellite ids to actions.
+_SCORE_EPS = 1e-9  # keeps every score strictly inside (0, 1)
 
-    Valid actions:
-    - capture
-    - downlink
-    - maintain
-    - idle
 
-    Decision guidance:
-    - 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
+def _clamp_score(value: float) -> float:
+    """Clamp *value* to the open interval (0, 1) exclusive."""
+    return max(_SCORE_EPS, min(1.0 - _SCORE_EPS, float(value)))
+
+# ── 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"
 
-    Output format:
-    {"0": "capture", "1": "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
 
-    Do not include explanations or any extra text outside the JSON object.
-    """
-).strip()
+    Output format (no explanation, no markdown):
+    {"0": "produce", "1": "assemble", "2": "deliver"}
+""").strip()
 
 
-# ── result dataclass ───────────────────────────────────────────────────────────
+# ── result dataclass ──────────────────────────────────────────────────────────
 @dataclass
 class TaskRunResult:
     task_name: str
-    grade: float                          # 0.0 – 1.0  (primary output)
-    grade_components: Dict[str, float]    # per-criterion, each 0.0 – 1.0
+    score: float
+    total_reward: float
     steps: int
     done: bool
-    metrics: Dict[str, float]             # raw env metrics (informational)
+    metrics: Dict[str, float]
 
 
 # ── observation helpers ────────────────────────────────────────────────────────
 
-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 {
-        "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]:
+def observation_to_dict(obs: ManufacturingObservation) -> Dict[str, Any]:
     return {
-        "satellites": [
+        "platforms": [
             {
-                "id": sat.id,
-                "position": sat.position,
-                "battery": sat.battery,
-                "storage": sat.storage,
-                "last_action": sat.last_action,
+                "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,
             }
-            for sat in observation.satellites
+            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
+        ],
+        "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
         ],
-        "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,
+        "total_reward": obs.total_reward,
+        "done": obs.done,
+        "reward": obs.reward,
+        "metadata": obs.metadata,
     }
 
 
-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
+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", [])}
 
 
 # ── heuristic policy ───────────────────────────────────────────────────────────
 
-def heuristic_action(observation: Dict[str, Any]) -> Dict[int, str]:
+def heuristic_action(obs_dict: Dict[str, Any]) -> Dict[int, str]:
     actions: Dict[int, str] = {}
-    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"
+    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"
 
+        actions[pid] = action
     return actions
 
 
-def safe_heuristic_action(observation: Dict[str, Any], reason: str) -> Dict[int, str]:
+def safe_heuristic_action(obs_dict: Dict[str, Any], reason: str) -> Dict[int, str]:
     try:
-        return heuristic_action(observation)
+        return heuristic_action(obs_dict)
     except Exception as exc:  # noqa: BLE001
         warn_once(
             f"heuristic:{reason}",
-            f"Heuristic fallback failed after {reason}: {exc}. Returning all-idle actions.",
+            f"Heuristic fallback failed after {reason}: {exc}. Returning all-{FALLBACK_ACTION}.",
         )
-        return build_idle_actions(observation)
+        return build_idle_actions(obs_dict)
 
 
 # ── prompt formatting ──────────────────────────────────────────────────────────
 
-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}"
+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']}"
         )
 
-    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()
+    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)"
+        )
+
+    solar_str = ", ".join(
+        f"{z}={round(v * 100)}%"
+        for z, v in obs_dict.get("solar_conditions", {}).items()
     )
 
-    return textwrap.dedent(
-        f"""
+    return textwrap.dedent(f"""
         Task: {task_name}
-        Time Step: {observation.get('time_step', 0)}
-        Total Reward: {observation.get('total_reward', 0.0):.2f}
-        Weather: {weather_info or 'n/a'}
+        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'}
 
-        Satellites:
-        {chr(10).join(satellites_info) or '  None'}
+        Open Delivery Windows:
+        {chr(10).join(windows_lines) if windows_lines else '  None'}
 
-        Pending Tasks:
-        {chr(10).join(tasks_info) or '  - None'}
-        """
-    ).strip()
+        Pending orders: {len(obs_dict.get('pending_orders', []))}
+    """).strip()
 
 
 def build_user_prompt(
     task_name: str,
     step: int,
-    observation: Dict[str, Any],
+    obs_dict: 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, observation)}
+        {format_observation(task_name, obs_dict)}
 
         Previous steps:
         {build_history_lines(history)}
 
         Reply with exactly one JSON object.
-        """
-    ).strip()
+    """).strip()
 
 
 # ── model response parsing ────────────────────────────────────────────────────
@@ -460,46 +318,43 @@ def extract_response_text(completion: Any) -> str:
     if isinstance(content, list):
         parts: List[str] = []
         for item in content:
-            if isinstance(item, dict):
-                text = item.get("text")
-            else:
-                text = getattr(item, "text", None)
+            text = item.get("text") if isinstance(item, dict) else 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, observation: Dict[str, Any]) -> Dict[int, str]:
+def parse_model_action(
+    response_text: str, obs_dict: Dict[str, Any]
+) -> Dict[int, str]:
     if not response_text:
-        return safe_heuristic_action(observation, "empty model response")
-
-    cleaned = ACTION_PREFIX_RE.sub("", response_text.strip())
+        return safe_heuristic_action(obs_dict, "empty model response")
 
     try:
-        json_match = re.search(r"\{.*\}", cleaned, re.DOTALL)
+        json_match = re.search(r"\{.*\}", response_text.strip(), 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():
-                sat_id = int(key)
-                if sat_id not in valid_ids:
+                pid    = int(key)
+                if pid not in valid_ids:
                     continue
                 action = str(value).strip().lower()
-                if not ACTION_PATTERN.fullmatch(action):
+                if action not in VALID_ACTIONS:
                     action = FALLBACK_ACTION
-                actions[sat_id] = action
+                actions[pid] = action
             if actions:
-                fallback = heuristic_action(observation)
-                for sat_id in valid_ids:
-                    actions.setdefault(sat_id, fallback.get(sat_id, FALLBACK_ACTION))
+                fallback = heuristic_action(obs_dict)
+                for pid in valid_ids:
+                    actions.setdefault(pid, fallback.get(pid, FALLBACK_ACTION))
                 return actions
     except (json.JSONDecodeError, TypeError, ValueError):
         pass
 
-    warn_once("parse:model-response", "Model response was not valid JSON; using heuristic actions.")
-    return safe_heuristic_action(observation, "invalid model response")
+    warn_once("parse:model-response", "Model response was not valid JSON; using heuristic.")
+    return safe_heuristic_action(obs_dict, "invalid model response")
 
 
 # ── client construction ────────────────────────────────────────────────────────
@@ -508,7 +363,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",
@@ -518,28 +373,47 @@ def validate_api_base_url(base_url: Optional[str]) -> Optional[str]:
     return cleaned
 
 
-def build_client() -> OpenAI:
-    return OpenAI(
-        base_url=API_BASE_URL,
-        api_key=API_KEY,
-        timeout=REQUEST_TIMEOUT,
-    )
+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
 
 
 # ── action chooser ─────────────────────────────────────────────────────────────
 
 async def choose_actions(
-    client: Optional[OpenAI],
+    client: Optional[Any],
     task_name: str,
     step: int,
-    observation: Dict[str, Any],
+    obs_dict: Dict[str, Any],
     history: List[str],
     total_reward: float,
 ) -> Dict[int, str]:
     if client is None:
-        raise RuntimeError("OpenAI client not initialized")
+        return safe_heuristic_action(obs_dict, "heuristic mode")
 
-    user_prompt = build_user_prompt(task_name, step, observation, history, total_reward)
+    user_prompt = build_user_prompt(task_name, step, obs_dict, history, total_reward)
 
     def _call() -> str:
         completion = client.chat.completions.create(
@@ -562,84 +436,81 @@ async def choose_actions(
     except asyncio.TimeoutError:
         warn_once(
             f"timeout:{task_name}",
-            f"[{task_name}] Step {step} timed out after {STEP_TIMEOUT}s. Using heuristic actions.",
+            f"[{task_name}] Step {step} timed out after {STEP_TIMEOUT}s. Using heuristic.",
         )
-        return safe_heuristic_action(observation, f"step timeout on {task_name} step {step}")
+        return safe_heuristic_action(obs_dict, 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 actions.",
+            f"[{task_name}] Model request failed at step {step}: {exc}. Using heuristic.",
         )
-        return safe_heuristic_action(observation, f"model request failure on {task_name} step {step}")
+        return safe_heuristic_action(obs_dict, f"model failure on {task_name} step {step}")
 
     if DEBUG:
-        print(f"[{task_name}] model response: {response_text[:300]}")
+        print(f"[DEBUG] [{task_name}] step={step} model_response={response_text[:300]!r}")
 
     try:
-        return parse_model_action(response_text, observation)
+        return parse_model_action(response_text, obs_dict)
     except Exception as exc:  # noqa: BLE001
         warn_once(
             f"parse:{task_name}",
-            f"[{task_name}] Failed to parse model response at step {step}: {exc}. "
-            "Using heuristic actions.",
+            f"[{task_name}] Parse failed at step {step}: {exc}. Using heuristic.",
         )
-        return safe_heuristic_action(observation, f"parse failure on {task_name} step {step}")
+        return safe_heuristic_action(obs_dict, f"parse failure on {task_name} step {step}")
 
 
 # ── episode runner ─────────────────────────────────────────────────────────────
 
-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)
+async def run_task(task_name: str, client: Optional[Any]) -> TaskRunResult:
+    env    = ManufacturingTaskEnv(task_name=task_name)
+    grader = ManufacturingTaskGrader(task_name=task_name)
     history: List[str] = []
 
-    observation = env.reset()
-    state = env.state()
+    obs        = env.reset()
+    state      = env.state()
     step_limit = state.max_steps
-    print(f"[START] task={task_name}", flush=True)
+
+    print(f"[START] task={task_name} max_steps={step_limit}", flush=True)
 
     for step in range(1, step_limit + 1):
-        obs_dict = observation_to_dict(observation)
-        actions = await choose_actions(
-            client,
-            task_name,
-            step,
-            obs_dict,
-            history,
-            observation.total_reward,
-        )
+        obs_dict = observation_to_dict(obs)
+        actions  = await 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)
 
-        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] step={step} reward={reward_value:.4f} total={observation.total_reward:.4f}",
+            f"[STEP]  task={task_name} step={step}/{step_limit}"
+            f" reward={reward_value:.4f} total={obs.total_reward:.4f}"
+            f" done={done}",
             flush=True,
         )
 
-        if REQUEST_DELAY > 0 and step < step_limit and not done:
-            time.sleep(REQUEST_DELAY)
+        if REQUEST_DELAY > 0 and not done:
+            await asyncio.sleep(REQUEST_DELAY)
 
         if done:
             break
 
     final_state = env.state()
-    metrics = {key: float(value) for key, value in final_state.metrics.items()}
-    grade = grader.grade_episode(env)
-    grade_components = grader.grade_components(env)
+    metrics     = {k: float(v) for k, v in final_state.metrics.items()}
+    score       = _clamp_score(grader.grade(metrics, final_state.step_count, final_state.platforms))
+
     print(
-        f"[END] task={task_name} grade={grade:.4f} "
-        f"steps={final_state.step_count} done={final_state.done}",
+        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}",
         flush=True,
     )
+
     return TaskRunResult(
         task_name=task_name,
-        grade=grade,
-        grade_components=grade_components,
+        score=score,
+        total_reward=final_state.total_reward,
         steps=final_state.step_count,
         done=final_state.done,
         metrics=metrics,
@@ -649,52 +520,30 @@ async def run_task(task_name: str, client: Optional[OpenAI]) -> TaskRunResult:
 # ── summary printer ────────────────────────────────────────────────────────────
 
 def print_summary(results: List[TaskRunResult]) -> None:
-    aggregate = sum(r.grade for r in results) / len(results)
-    all_keys: list = []
-    for r in results:
-        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)
+    aggregate = sum(r.score for r in results) / len(results)
+    print("\nInference Summary")
+    print("=" * 60)
     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)
+        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)
 
 
 # ── async main ─────────────────────────────────────────────────────────────────
 
 async def async_main() -> None:
-    client = build_client()
+    client  = build_client()
     results = []
     for task_name in TASK_ORDER:
         if TASK_TIMEOUT > 0:
             try:
-                result = await asyncio.wait_for(
-                    run_task(task_name, client), timeout=TASK_TIMEOUT
-                )
+                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.",