fix: zero-arg self-contained graders, inference runs all 3 tasks

grader.py

@@ -1,16 +1,21 @@
 """
 Graders for Smart Factory Scheduling tasks.
-Called by the OpenEnv validator to score an episode.
 
-Each grader accepts either:
-  - an env object  (has .completed_jobs, .jobs, .time, .late_jobs attributes)
-  - a state dict   (has "completed_jobs", "pending_jobs", "time", "late_jobs" keys)
+Each grader is self-contained: when called with no arguments it creates a
+FactoryEnv, runs a deterministic heuristic episode, and returns a score
+strictly in (0, 1).
 
-Returns a float strictly in (0, 1).
+Alternatively, pass an env object or state dict from an already-run episode:
+    score_easy(env)       # env object with .completed_jobs, .jobs, .time …
+    score_easy(state)     # dict with "completed_jobs", "pending_jobs", "time" …
 """
 
+from __future__ import annotations
 
-def _compute(completed, on_time, total, late):
+
+# ── internal helpers ──────────────────────────────────────────────────────────
+
+def _compute(completed: int, on_time: int, total: int, late: int) -> float:
     if total == 0:
         return 0.001
     score = (
@@ -21,37 +26,79 @@ def _compute(completed, on_time, total, late):
     return round(max(0.001, min(0.999, score)), 4)
 
 
-def _score(state_or_env):
-    if isinstance(state_or_env, dict):
-        done = state_or_env.get("completed_jobs", []) or []
-        pending = state_or_env.get("pending_jobs", []) or []
-        late = state_or_env.get("late_jobs", 0) or 0
-        t = state_or_env.get("time", 0) or 0
+def _score_from(obj) -> float:
+    """Accept env object or state dict and return a score."""
+    if isinstance(obj, dict):
+        done_jobs = obj.get("completed_jobs", []) or []
+        pending   = obj.get("pending_jobs", []) or []
+        late      = obj.get("late_jobs", 0) or 0
+        t         = obj.get("time", 0) or 0
     else:
-        done = list(getattr(state_or_env, "completed_jobs", []) or [])
-        pending = list(getattr(state_or_env, "jobs", getattr(state_or_env, "pending_jobs", [])) or [])
-        late = getattr(state_or_env, "late_jobs", 0) or 0
-        t = getattr(state_or_env, "time", 0) or 0
-
-    completed = len(done)
-    total = completed + len(pending)
-    on_time = sum(
-        1 for j in done
-        if (j.get("deadline", 0) if isinstance(j, dict) else getattr(j, "deadline", 0)) >= t
+        done_jobs = list(getattr(obj, "completed_jobs", []) or [])
+        pending   = list(
+            getattr(obj, "jobs", getattr(obj, "pending_jobs", []))
+        ) or []
+        late = getattr(obj, "late_jobs", 0) or 0
+        t    = getattr(obj, "time", 0) or 0
+
+    completed = len(done_jobs)
+    total     = completed + len(pending)
+    on_time   = sum(
+        1 for j in done_jobs
+        if (j.get("deadline", 0) if isinstance(j, dict)
+            else getattr(j, "deadline", 0)) >= t
     )
     return _compute(completed, on_time, total, late)
 
 
-def score_easy(state_or_env):
+def _heuristic_action(obs):
+    """Simple earliest-deadline-first heuristic."""
+    from factory_env.models import FactoryAction
+    for m in obs.machines:
+        if m.status == "broken":
+            return FactoryAction(action_type="repair", machine_id=m.id)
+    for j in sorted(obs.pending_jobs, key=lambda x: (x.deadline, -x.priority)):
+        for m in obs.machines:
+            if m.status == "idle":
+                return FactoryAction(
+                    action_type="assign_job", job_id=j.id, machine_id=m.id
+                )
+    return None  # wait
+
+
+def _run_episode(task: str, seed: int = 42) -> float:
+    """Run one full heuristic episode and return the graded score."""
+    from factory_env.env import FactoryEnv
+    from factory_env.models import FactoryAction
+
+    env = FactoryEnv(task=task, seed=seed)
+    obs = env.reset()
+    for _ in range(obs.max_steps):
+        if obs.done:
+            break
+        action = _heuristic_action(obs) or FactoryAction(action_type="wait")
+        obs = env.step(action)
+    return _score_from(env)
+
+
+# ── public graders ────────────────────────────────────────────────────────────
+
+def score_easy(state_or_env=None) -> float:
     """Grade an easy-task episode (2 machines, 3 jobs, no failures)."""
-    return _score(state_or_env)
+    if state_or_env is not None:
+        return _score_from(state_or_env)
+    return _run_episode("easy")
 
 
-def score_medium(state_or_env):
+def score_medium(state_or_env=None) -> float:
     """Grade a medium-task episode (4 machines, 7 jobs, 8% failures)."""
-    return _score(state_or_env)
+    if state_or_env is not None:
+        return _score_from(state_or_env)
+    return _run_episode("medium")
 
 
-def score_hard(state_or_env):
+def score_hard(state_or_env=None) -> float:
     """Grade a hard-task episode (6 machines, 12 jobs, 15% failures)."""
-    return _score(state_or_env)
+    if state_or_env is not None:
+        return _score_from(state_or_env)
+    return _run_episode("hard")

inference.py

@@ -1,20 +1,19 @@
 """
 Inference Script — Smart Factory Scheduling Environment
 =======================================================
-Connects to a running factory_env server via WebSocket and runs an LLM agent.
-
-Mandatory env vars (per hackathon spec):
-  HF_TOKEN       HuggingFace API key (also used as OPENAI_API_KEY)
-  API_BASE_URL   LLM endpoint  (default: HF router)
-  MODEL_NAME     Model ID      (default: Qwen/Qwen2.5-72B-Instruct)
-
-Optional env vars:
-  ENV_URL        URL of running factory_env server (default: http://localhost:7860)
-  IMAGE_NAME     Docker image name — if set, spins up a container instead of ENV_URL
-  FACTORY_TASK   easy | medium | hard  (default: easy)
-
-STDOUT FORMAT (strict — do not alter):
-  [START] task=<name> env=factory_env model=<model>
+Runs an LLM agent against the factory_env server for all 3 tasks
+(easy, medium, hard) and emits structured stdout logs.
+
+Environment variables:
+  HF_TOKEN        HuggingFace / API key  (no default — required)
+  API_BASE_URL    LLM endpoint           (default: HF router)
+  MODEL_NAME      Model identifier       (default: Qwen/Qwen2.5-72B-Instruct)
+  IMAGE_NAME      Docker image name — if set, spins up a container
+  ENV_URL         Server URL             (default: http://localhost:7860)
+  FACTORY_TASK    Run a single task: easy | medium | hard  (default: run all 3)
+
+STDOUT FORMAT  (one [START] / N [STEP] / one [END] per task):
+  [START] task=<task> env=factory_env model=<model>
   [STEP]  step=<n> action=<str> reward=<0.00> done=<true|false> error=<msg|null>
   [END]   success=<true|false> steps=<n> score=<0.000> rewards=<r1,r2,...>
 """
@@ -32,15 +31,19 @@ from factory_env.models import FactoryAction
 
 # ── Configuration ────────────────────────────────────────────────────────────
 HF_TOKEN = os.getenv("HF_TOKEN")
+API_KEY   = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
 API_BASE_URL: str = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")
-MODEL_NAME: str = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct")
-TASK_NAME: str = os.getenv("FACTORY_TASK", "easy")
-ENV_URL: str = os.getenv("ENV_URL", "http://localhost:7860")
-LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")
-BENCHMARK: str = "factory_env"
-TEMPERATURE: float = 0.2
-MAX_TOKENS: int = 80
-SUCCESS_SCORE_THRESHOLD: float = 0.5
+MODEL_NAME: str   = os.getenv("MODEL_NAME",   "Qwen/Qwen2.5-72B-Instruct")
+IMAGE_NAME        = os.getenv("IMAGE_NAME") or os.getenv("LOCAL_IMAGE_NAME")
+ENV_URL: str      = os.getenv("ENV_URL", "http://localhost:7860")
+BENCHMARK: str    = "factory_env"
+TEMPERATURE: float        = 0.2
+MAX_TOKENS: int           = 80
+SUCCESS_SCORE_THRESHOLD   = 0.5
+
+# Run a single task if FACTORY_TASK is set, otherwise run all three
+_single = os.getenv("FACTORY_TASK", "").strip()
+TASKS: List[str] = [_single] if _single else ["easy", "medium", "hard"]
 
 SYSTEM_PROMPT = textwrap.dedent("""
     You are controlling a smart factory scheduling system.
@@ -53,14 +56,13 @@ SYSTEM_PROMPT = textwrap.dedent("""
 """).strip()
 
 
-# ── Log helpers (strict format required by judges) ────────────────────────────
+# ── Log helpers ───────────────────────────────────────────────────────────────
 def log_start(task: str, env: str, model: str) -> None:
     print(f"[START] task={task} env={env} model={model}", flush=True)
 
 
-def log_step(
-    step: int, action: str, reward: float, done: bool, error: Optional[str]
-) -> None:
+def log_step(step: int, action: str, reward: float, done: bool,
+             error: Optional[str]) -> None:
     print(
         f"[STEP] step={step} action={action.replace(' ', '_')} reward={reward:.2f} "
         f"done={str(done).lower()} error={error or 'null'}",
@@ -68,9 +70,8 @@ def log_step(
     )
 
 
-def log_end(
-    success: bool, steps: int, score: float, rewards: List[float]
-) -> None:
+def log_end(success: bool, steps: int, score: float,
+            rewards: List[float]) -> None:
     print(
         f"[END] success={str(success).lower()} steps={steps} "
         f"score={score:.3f} rewards={','.join(f'{r:.2f}' for r in rewards)}",
@@ -86,7 +87,8 @@ def build_prompt(step: int, obs, last_reward: float) -> str:
     )
     jobs = (
         "\n".join(
-            f"  {j.id}: remaining={j.remaining_time}, deadline={j.deadline}, priority={j.priority}"
+            f"  {j.id}: remaining={j.remaining_time}, deadline={j.deadline},"
+            f" priority={j.priority}"
             for j in obs.pending_jobs
         )
         or "  (none)"
@@ -104,7 +106,7 @@ def get_model_action(client: OpenAI, step: int, obs, last_reward: float) -> str:
             model=MODEL_NAME,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": build_prompt(step, obs, last_reward)},
+                {"role": "user",   "content": build_prompt(step, obs, last_reward)},
             ],
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
@@ -122,7 +124,8 @@ def parse_action(text: str) -> FactoryAction:
     try:
         parts = text.strip().split()
         if parts[0] == "assign_job" and len(parts) == 3:
-            return FactoryAction(action_type="assign_job", job_id=parts[1], machine_id=parts[2])
+            return FactoryAction(action_type="assign_job",
+                                 job_id=parts[1], machine_id=parts[2])
         if parts[0] == "repair" and len(parts) == 2:
             return FactoryAction(action_type="repair", machine_id=parts[1])
     except Exception:
@@ -131,7 +134,6 @@ def parse_action(text: str) -> FactoryAction:
 
 
 def heuristic_action(obs):
-    """Fallback heuristic when LLM returns an ineffective wait."""
     for m in obs.machines:
         if m.status == "broken":
             return FactoryAction(action_type="repair", machine_id=m.id), f"repair {m.id}"
@@ -139,54 +141,41 @@ def heuristic_action(obs):
         for m in obs.machines:
             if m.status == "idle":
                 s = f"assign_job {j.id} {m.id}"
-                return FactoryAction(action_type="assign_job", job_id=j.id, machine_id=m.id), s
+                return FactoryAction(action_type="assign_job",
+                                     job_id=j.id, machine_id=m.id), s
     return FactoryAction(action_type="wait"), "wait"
 
 
 # ── Score from final state ────────────────────────────────────────────────────
 def score_from_state(state, task: str) -> float:
-    """Compute episode score from WebSocket state response."""
     completed_jobs = getattr(state, "completed_jobs", []) or []
-    pending_jobs = getattr(state, "pending_jobs", []) or []
-    late_jobs = getattr(state, "late_jobs", 0) or 0
-    time = getattr(state, "time", 0) or 0
-
-    completed = len(completed_jobs)
-    total = completed + len(pending_jobs)
-
-    # on_time: jobs whose deadline hasn't passed by end of episode (matches grader)
+    pending_jobs   = getattr(state, "pending_jobs",   []) or []
+    late_jobs      = getattr(state, "late_jobs", 0) or 0
+    time           = getattr(state, "time",      0) or 0
+    completed      = len(completed_jobs)
+    total          = completed + len(pending_jobs)
     on_time = sum(
         1 for j in completed_jobs
-        if (j.get("deadline", 0) if isinstance(j, dict) else j.deadline) >= time
+        if (j.get("deadline", 0) if isinstance(j, dict)
+            else j.deadline) >= time
     )
-
     return compute_score(completed, on_time, total, late_jobs, task)
 
 
-# ── Main ──────────────────────────────────────────────────────────────────────
-async def main() -> None:
-    llm_client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
-
+# ── Single-task episode ───────────────────────────────────────────────────────
+async def run_task(env_client: FactoryEnvClient,
+                   llm_client: OpenAI,
+                   task: str) -> None:
     rewards: List[float] = []
     steps_taken = 0
-    score = 0.0
-    success = False
-
-    log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
+    score       = 0.0
+    success     = False
 
-    # Connect to environment — Docker image, direct URL, or localhost
-    if LOCAL_IMAGE_NAME:
-        print(f"[DEBUG] Spinning up Docker image: {LOCAL_IMAGE_NAME}", flush=True)
-        env = await FactoryEnvClient.from_docker_image(LOCAL_IMAGE_NAME)
-    else:
-        url = ENV_URL or "http://localhost:7860"
-        print(f"[DEBUG] Connecting to: {url}", flush=True)
-        env = FactoryEnvClient(base_url=url)
-        await env.connect()
+    log_start(task=task, env=BENCHMARK, model=MODEL_NAME)
 
     try:
-        result = await env.reset(task=TASK_NAME)
-        obs = result.observation
+        result     = await env_client.reset(task=task)
+        obs        = result.observation
         last_reward = 0.0
 
         for step in range(1, obs.max_steps + 1):
@@ -194,44 +183,64 @@ async def main() -> None:
                 break
 
             action_text = get_model_action(llm_client, step, obs, last_reward)
-            action = parse_action(action_text)
+            action      = parse_action(action_text)
 
-            # Heuristic fallback: if LLM returns wait but there's work to do
             if action.action_type == "wait" and (
-                obs.pending_jobs or any(m.status == "broken" for m in obs.machines)
+                obs.pending_jobs
+                or any(m.status == "broken" for m in obs.machines)
             ):
                 action, action_text = heuristic_action(obs)
 
-            result = await env.step(action)
-            obs = result.observation
-            reward = result.reward or 0.0
-            done = result.done
+            result      = await env_client.step(action)
+            obs         = result.observation
+            reward      = result.reward or 0.0
+            done        = result.done
             rewards.append(reward)
             steps_taken = step
             last_reward = reward
 
-            log_step(step=step, action=action_text, reward=reward, done=done, error=None)
-
+            log_step(step=step, action=action_text,
+                     reward=reward, done=done, error=None)
             if done:
                 break
 
-        # Compute score from final WebSocket state
         try:
-            state = await env.state()
-            score = score_from_state(state, TASK_NAME)
+            state = await env_client.state()
+            score = score_from_state(state, task)
         except Exception as exc:
-            print(f"[DEBUG] state() failed, falling back to reward sum: {exc}", flush=True)
-            max_reward = {"easy": 4.0, "medium": 12.0, "hard": 20.0}.get(TASK_NAME, 10.0)
-            score = min(max(sum(rewards) / max_reward, 0.0), 1.0)
+            print(f"[DEBUG] state() failed: {exc}", flush=True)
+            max_r = {"easy": 4.0, "medium": 12.0, "hard": 20.0}.get(task, 10.0)
+            raw   = sum(rewards) / max_r if max_r > 0 else 0.0
+            score = round(max(0.001, min(0.999, raw)), 4)
 
         success = score >= SUCCESS_SCORE_THRESHOLD
 
+    finally:
+        log_end(success=success, steps=steps_taken,
+                score=score, rewards=rewards)
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+async def main() -> None:
+    llm_client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+
+    if IMAGE_NAME:
+        print(f"[DEBUG] Spinning up Docker image: {IMAGE_NAME}", flush=True)
+        env_client = await FactoryEnvClient.from_docker_image(IMAGE_NAME)
+    else:
+        url = ENV_URL or "http://localhost:7860"
+        print(f"[DEBUG] Connecting to: {url}", flush=True)
+        env_client = FactoryEnvClient(base_url=url)
+        await env_client.connect()
+
+    try:
+        for task in TASKS:
+            await run_task(env_client, llm_client, task)
     finally:
         try:
-            await env.close()
+            await env_client.close()
         except Exception as exc:
             print(f"[DEBUG] env.close() error: {exc}", flush=True)
-        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
 
 
 if __name__ == "__main__":