feat: init inference

inference.py

@@ -1,48 +1,19 @@
 """
-Inference Script Example
+Inference Script — WhyDidItFail
 ===================================
-MANDATORY
-- Before submitting, ensure the following variables are defined in your environment configuration:
-    API_BASE_URL   The API endpoint for the LLM.
-    MODEL_NAME     The model identifier to use for inference.
-    HF_TOKEN       Your Hugging Face / API key.
-    LOCAL_IMAGE_NAME The name of the local image to use for the environment if you are using from_docker_image()
-                     method
-
-- Defaults are set only for API_BASE_URL and MODEL_NAME 
-    (and should reflect your active inference setup):
-    API_BASE_URL = os.getenv("API_BASE_URL", "<your-active-endpoint>")
-    MODEL_NAME = os.getenv("MODEL_NAME", "<your-active-model>")
-    
-- The inference script must be named `inference.py` and placed in the root directory of the project
-- Participants must use OpenAI Client for all LLM calls using above variables
+MANDATORY environment variables:
+    API_BASE_URL        The API endpoint for the LLM.
+    MODEL_NAME          The model identifier to use for inference.
+    HF_TOKEN / API_KEY  Your Hugging Face / API key.
 
 STDOUT FORMAT
-- The script must emit exactly three line types to stdout, in this order:
-
     [START] task=<task_name> env=<benchmark> model=<model_name>
     [STEP]  step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>
     [END]   success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>
-
-  Rules:
-    - One [START] line at episode begin.
-    - One [STEP] line per step, immediately after env.step() returns.
-    - One [END] line after env.close(), always emitted (even on exception).
-    - reward and rewards are formatted to 2 decimal places.
-    - done and success are lowercase booleans: true or false.
-    - error is the raw last_action_error string, or null if none.
-    - All fields on a single line with no newlines within a line.
-    - Each tasks should return score in [0, 1]
-
-  Example:
-    [START] task=click-test env=miniwob model=Qwen3-VL-30B
-    [STEP] step=1 action=click('123') reward=0.00 done=false error=null
-    [STEP] step=2 action=fill('456','text') reward=0.00 done=false error=null
-    [STEP] step=3 action=click('789') reward=1.00 done=true error=null
-    [END] success=true steps=3 score=1.00 rewards=0.00,0.00,1.00
 """
 
 import asyncio
+import json
 import os
 import textwrap
 from typing import List, Optional
@@ -52,31 +23,38 @@ load_dotenv()
 
 from openai import OpenAI
 
-from WhyDidItFail.client import WhydiditfailEnv
-from WhyDidItFail.models import WhyDidItFailAction
-IMAGE_NAME = os.getenv("IMAGE_NAME") # If you are using docker image
-API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+from client import WhyDidItFailEnv
+from models import WhyDidItFailAction
 
+IMAGE_NAME = os.getenv("IMAGE_NAME")
+API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
 API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1"
 MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct"
 TASK_NAME = os.getenv("WHYDIDITFAIL_TASK", "whydiditfail")
 BENCHMARK = os.getenv("WHYDIDITFAIL_BENCHMARK", "whydiditfail")
 MAX_STEPS = 8
-TEMPERATURE = 0.7
-MAX_TOKENS = 150
-SUCCESS_SCORE_THRESHOLD = 0.1  # normalized score in [0, 1]
-
-# Max possible reward: each token contributes 0.1, across all steps
-_MAX_REWARD_PER_STEP = MAX_TOKENS * 0.1
-MAX_TOTAL_REWARD = MAX_STEPS * _MAX_REWARD_PER_STEP
+TEMPERATURE = 0.3
+MAX_TOKENS = 256
+SUCCESS_SCORE_THRESHOLD = 0.5  # reward >= 0.5 counts as success
 
 SYSTEM_PROMPT = textwrap.dedent(
     """
-    You are interacting with a simple echo environment.
-    Each turn you must send a message. The environment will echo it back.
-    Reward is proportional to message length: reward = len(message) * 0.1
-    Your goal is to maximize total reward by sending meaningful, substantive messages.
-    Reply with exactly one message string — no quotes, no prefixes, just the message text.
+    You are a machine learning engineer diagnosing a failed training run.
+    Each turn you will receive data from the training run and must decide what to investigate next.
+
+    Available actions:
+    - inspect_logs       : examine training loss curves
+    - inspect_config     : examine hyperparameter config (lr, optimizer, etc.)
+    - inspect_gradients  : examine gradient statistics
+    - submit_diagnosis   : submit your final diagnosis (ends the episode)
+
+    You must respond with a JSON object on a single line. Examples:
+        {"action_type": "inspect_logs"}
+        {"action_type": "inspect_config"}
+        {"action_type": "submit_diagnosis", "diagnosis": "exploding gradients"}
+
+    Only submit_diagnosis when you are confident. The diagnosis should describe the failure mode
+    in plain terms (e.g. "exploding gradients", "overfitting", "vanishing gradients").
     """
 ).strip()
 
@@ -87,11 +65,7 @@ def log_start(task: str, env: str, model: str) -> None:
 
 def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
     error_val = error if error else "null"
-    done_val = str(done).lower()
-    print(
-        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
-        flush=True,
-    )
+    print(f"[STEP] step={step} action={action} reward={reward:.2f} done={str(done).lower()} error={error_val}", flush=True)
 
 
 def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
@@ -99,22 +73,25 @@ def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> No
     print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True)
 
 
-def build_user_prompt(step: int, last_echoed: str, last_reward: float, history: List[str]) -> str:
+def build_user_prompt(step: int, observation_summary: str, history: List[str]) -> str:
     history_block = "\n".join(history[-4:]) if history else "None"
     return textwrap.dedent(
         f"""
         Step: {step}
-        Last echoed message: {last_echoed!r}
-        Last reward: {last_reward:.2f}
-        Previous steps:
+
+        Current observation:
+        {observation_summary}
+
+        History:
         {history_block}
-        Send your next message.
+
+        Respond with a JSON action.
         """
     ).strip()
 
 
-def get_model_message(client: OpenAI, step: int, last_echoed: str, last_reward: float, history: List[str]) -> str:
-    user_prompt = build_user_prompt(step, last_echoed, last_reward, history)
+def get_model_action(client: OpenAI, step: int, observation_summary: str, history: List[str]) -> WhyDidItFailAction:
+    user_prompt = build_user_prompt(step, observation_summary, history)
     try:
         completion = client.chat.completions.create(
             model=MODEL_NAME,
@@ -127,16 +104,30 @@ def get_model_message(client: OpenAI, step: int, last_echoed: str, last_reward:
             stream=False,
         )
         text = (completion.choices[0].message.content or "").strip()
-        return text if text else "hello"
+        data = json.loads(text)
+        return WhyDidItFailAction(**data)
     except Exception as exc:
-        print(f"[DEBUG] Model request failed: {exc}", flush=True)
-        return "hello"
+        print(f"[DEBUG] Model request/parse failed: {exc}", flush=True)
+        # Fallback: inspect logs if early, otherwise give up and submit empty diagnosis
+        if step <= 2:
+            return WhyDidItFailAction(action_type="inspect_logs")
+        return WhyDidItFailAction(action_type="submit_diagnosis", diagnosis="unknown")
+
+
+def summarize_observation(obs) -> str:
+    lines = [
+        f"Task: {obs.task_description}",
+        f"Feedback: {obs.feedback}",
+        f"Available actions: {', '.join(obs.available_actions)}",
+    ]
+    if obs.visible_data:
+        lines.append(f"Data: {json.dumps(obs.visible_data, indent=2)}")
+    return "\n".join(lines)
 
 
 async def main() -> None:
     client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
-
-    env = await WhydiditfailEnv.from_docker_image(IMAGE_NAME or "")
+    env = await WhyDidItFailEnv.from_docker_image(IMAGE_NAME or "")
 
     history: List[str] = []
     rewards: List[float] = []
@@ -147,44 +138,40 @@ async def main() -> None:
     log_start(task=TASK_NAME, env=BENCHMARK, model=MODEL_NAME)
 
     try:
-        result = await env.reset() # OpenENV.reset()
-        last_echoed = result.observation.echoed_message
-        last_reward = 0.0
+        result = await env.reset()
+        obs = result.observation
 
         for step in range(1, MAX_STEPS + 1):
             if result.done:
                 break
 
-            message = get_model_message(client, step, last_echoed, last_reward, history)
+            obs_summary = summarize_observation(obs)
+            action = get_model_action(client, step, obs_summary, history)
 
-            result = await env.step(WhyDidItFailAction(message=message))
+            result = await env.step(action)
             obs = result.observation
 
             reward = result.reward or 0.0
             done = result.done
-            error = None
+            action_str = action.model_dump_json(exclude_none=True)
 
             rewards.append(reward)
             steps_taken = step
-            last_echoed = obs.echoed_message
-            last_reward = reward
-
-            log_step(step=step, action=message, reward=reward, done=done, error=error)
 
-            history.append(f"Step {step}: {message!r} -> reward {reward:+.2f}")
+            log_step(step=step, action=action_str, reward=reward, done=done, error=None)
+            history.append(f"Step {step}: {action_str} -> reward={reward:.2f} feedback={obs.feedback!r}")
 
             if done:
                 break
 
-        score = sum(rewards) / MAX_TOTAL_REWARD if MAX_TOTAL_REWARD > 0 else 0.0
-        score = min(max(score, 0.0), 1.0)  # clamp to [0, 1]
+        score = max(rewards) if rewards else 0.0  # final diagnosis reward is what matters
         success = score >= SUCCESS_SCORE_THRESHOLD
 
     finally:
         try:
             await env.close()
         except Exception as e:
-            print(f"[DEBUG] env.close() error (container cleanup): {e}", flush=True)
+            print(f"[DEBUG] env.close() error: {e}", flush=True)
         log_end(success=success, steps=steps_taken, score=score, rewards=rewards)