fix: correct step format, per-task output, tuple grader returns

inference.py

@@ -1,16 +1,7 @@
 """
 Inference Script - ConfigDebugEnv
 ===================================
-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.
-    IMAGE_NAME     The name of the local image to use for the environment if you are using
-                   from_docker_image() method
-
-STDOUT FORMAT
-- The script emits [START], [STEP], and [END] PER TASK:
+STDOUT FORMAT - emits [START], [STEP], [END] PER TASK:
     [START] task=<task_id> 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>
@@ -28,14 +19,16 @@ MAX_STEPS_PER_TASK = 5
 TEMPERATURE = 0.1
 MAX_TOKENS = 2000
 
-SYSTEM_PROMPT = textwrap.dedent(
-    """
+SYSTEM_PROMPT = textwrap.dedent("""
     You are an expert DevOps engineer specializing in configuration file debugging.
     You will be given a broken configuration file and must fix ALL bugs in it.
     Return ONLY the fixed configuration file content.
     No explanations, no markdown formatting, no code blocks. Just the raw fixed configuration.
-    """
-).strip()
+""").strip()
+
+
+def clamp(v: float) -> float:
+    return max(0.01, min(0.99, v))
 
 
 def log_start(task: str, env: str, model: str) -> None:
@@ -43,22 +36,13 @@ 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()
-    reward = max(0.01, min(0.99, reward))
-    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={clamp(reward):.2f} done={str(done).lower()} error={error or 'null'}", flush=True)
 
 
 def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
-    score = max(0.01, min(0.99, score))
-    rewards_str = ",".join(f"{max(0.01, min(0.99, r)):.2f}" for r in rewards)
-    print(
-        f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",
-        flush=True,
-    )
+    s = clamp(score)
+    rs = ",".join(f"{clamp(r):.2f}" for r in rewards) if rewards else "0.01"
+    print(f"[END] success={str(success).lower()} steps={steps} score={s:.3f} rewards={rs}", flush=True)
 
 
 def strip_code_blocks(text: str) -> str:
@@ -73,38 +57,45 @@ def strip_code_blocks(text: str) -> str:
     return text
 
 
-def build_user_prompt(obs: dict, step: int, history: List[str]) -> str:
-    history_block = "\n".join(history[-4:]) if history else "None"
-    return textwrap.dedent(
-        f"""
-        Fix the following broken {obs.get('file_type', 'config')} configuration file.
+def get_obs_field(data: dict, field: str, default=None):
+    """Get field from response - handles both nested and flat formats."""
+    obs = data.get("observation", data)
+    return obs.get(field, data.get(field, default))
+
 
-        Task: {obs.get('task_description', '')}
-        Difficulty: {obs.get('difficulty', '')}
-        Number of bugs to find: {obs.get('num_bugs', 0)}
-        Bugs fixed so far: {obs.get('bugs_found_so_far', 0)}
-        Error message: {obs.get('error_message', '')}
-        Step: {step}
+def get_model_fix(client, obs_data, step, history, model_name):
+    file_type = get_obs_field(obs_data, "file_type", "config")
+    desc = get_obs_field(obs_data, "task_description", "")
+    difficulty = get_obs_field(obs_data, "difficulty", "")
+    num_bugs = get_obs_field(obs_data, "num_bugs", 0)
+    bugs_found = get_obs_field(obs_data, "bugs_found_so_far", 0)
+    error_msg = get_obs_field(obs_data, "error_message", "")
+    broken = get_obs_field(obs_data, "broken_config", "")
 
-        Previous attempts:
-        {history_block}
+    history_block = "\n".join(history[-4:]) if history else "None"
+    prompt = f"""Fix the following broken {file_type} configuration file.
+
+Task: {desc}
+Difficulty: {difficulty}
+Number of bugs to find: {num_bugs}
+Bugs fixed so far: {bugs_found}
+Error message: {error_msg}
+Step: {step}
 
-        Broken configuration:
-    {obs.get('broken_config', '')}
+Previous attempts:
+{history_block}
 
-        Return ONLY the fixed configuration file content.
-        """
-    ).strip()
+Broken configuration:
+{broken}
 
+Return ONLY the fixed configuration file content."""
 
-def get_model_fix(client: OpenAI, obs: dict, step: int, history: List[str], model_name: str) -> str:
-    user_prompt = build_user_prompt(obs, step, history)
     try:
         completion = client.chat.completions.create(
             model=model_name,
             messages=[
                 {"role": "system", "content": SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
+                {"role": "user", "content": prompt},
             ],
             temperature=TEMPERATURE,
             max_tokens=MAX_TOKENS,
@@ -112,133 +103,106 @@ def get_model_fix(client: OpenAI, obs: dict, step: int, history: List[str], mode
         )
         text = (completion.choices[0].message.content or "").strip()
         return strip_code_blocks(text) if text else ""
-    except Exception as exc:
-        print(f"[DEBUG] Model request failed: {exc}", flush=True)
+    except Exception as e:
+        print(f"[DEBUG] LLM error: {e}", flush=True)
         return ""
 
 
 class HTTPEnvClient:
-    def __init__(self, base_url: str):
+    def __init__(self, base_url):
         import httpx
         self.base_url = base_url.rstrip("/")
         self.http = httpx.AsyncClient(timeout=60.0)
 
     async def reset(self):
-        resp = await self.http.post(f"{self.base_url}/reset")
-        resp.raise_for_status()
-        return resp.json()
-
-    async def step(self, action_data: dict):
-        resp = await self.http.post(f"{self.base_url}/step", json=action_data)
-        resp.raise_for_status()
-        return resp.json()
+        r = await self.http.post(f"{self.base_url}/reset")
+        r.raise_for_status()
+        return r.json()
+
+    async def step(self, fixed_config: str):
+        """Send step with correct OpenEnv format: {"action": {"fixed_config": "..."}}"""
+        r = await self.http.post(
+            f"{self.base_url}/step",
+            json={"action": {"fixed_config": fixed_config}},
+        )
+        r.raise_for_status()
+        return r.json()
 
     async def close(self):
         await self.http.aclose()
 
 
-async def main() -> None:
+async def main():
     api_key = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or ""
     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"
-    benchmark = os.getenv("CONFIG_DEBUG_BENCHMARK", "config_debug_env")
+    benchmark = "config_debug_env"
 
     client = OpenAI(base_url=api_base_url, api_key=api_key)
-
     env_url = sys.argv[1] if len(sys.argv) > 1 else "http://localhost:7860"
     env = HTTPEnvClient(env_url)
 
     try:
-        # Reset environment — starts at task1
         result = await env.reset()
-        obs = result.get("observation", {})
-        state = result.get("state", obs)
-
-        current_task_id = obs.get("task_id", "unknown")
+        current_task = get_obs_field(result, "task_id", "unknown")
         task_step = 0
-        task_rewards: List[float] = []
-        task_history: List[str] = []
-
-        # Emit [START] for the first task
-        log_start(task=current_task_id, env=benchmark, model=model_name)
-
-        while True:
-            # Check if episode is fully done
-            is_done = state.get("is_done", False) if isinstance(state, dict) else False
-            if is_done:
-                # End the current task
-                if task_rewards:
-                    task_score = sum(task_rewards) / len(task_rewards)
-                else:
-                    task_score = 0.01
-                task_score = max(0.01, min(0.99, task_score))
-                log_end(
-                    success=task_score >= 0.5,
-                    steps=task_step,
-                    score=task_score,
-                    rewards=task_rewards if task_rewards else [0.01],
-                )
-                break
+        task_rewards = []
+        task_history = []
 
+        log_start(task=current_task, env=benchmark, model=model_name)
+
+        for global_step in range(1, 50):
+            fixed_config = get_model_fix(client, result, task_step + 1, task_history, model_name)
             task_step += 1
 
-            # Get LLM fix
-            fixed_config = get_model_fix(client, obs, task_step, task_history, model_name)
-
-            # Step the environment
-            step_result = await env.step({"fixed_config": fixed_config})
-            obs = step_result.get("observation", {})
-            state = step_result.get("state", obs)
-            reward = step_result.get("reward", 0.01)
-            reward = max(0.01, min(0.99, float(reward) if reward is not None else 0.01))
-            info = step_result.get("info", {})
-            new_task_id = obs.get("task_id", current_task_id)
-            task_done = info.get("task_done", False)
-            is_done = state.get("is_done", False) if isinstance(state, dict) else False
-            error = info.get("error_message") if info.get("error_message") != "All checks passed!" else None
+            try:
+                step_result = await env.step(fixed_config)
+            except Exception as step_err:
+                print(f"[DEBUG] Step failed: {step_err}", flush=True)
+                log_step(task_step, f"fix({current_task})", 0.01, False, str(step_err))
+                task_rewards.append(0.01)
+                # End this task on step failure
+                log_end(False, task_step, 0.01, task_rewards)
+                break
 
+            reward = get_obs_field(step_result, "reward", 0.01)
+            if reward is None:
+                reward = 0.01
+            reward = clamp(float(reward))
             task_rewards.append(reward)
 
-            action_summary = f"fix({current_task_id})"
-            log_step(
-                step=task_step,
-                action=action_summary,
-                reward=reward,
-                done=is_done,
-                error=error,
-            )
+            new_task = get_obs_field(step_result, "task_id", current_task)
+            is_done = get_obs_field(step_result, "done", False)
+            error = get_obs_field(step_result, "error_message", None)
+            if error == "All checks passed!":
+                error = None
 
+            log_step(task_step, f"fix({current_task})", reward, bool(is_done), error)
             task_history.append(f"Step {task_step}: reward {reward:.2f}")
 
-            # Did this task just finish?
-            if task_done or new_task_id != current_task_id:
-                # Emit [END] for the completed task
-                if task_rewards:
-                    task_score = sum(task_rewards) / len(task_rewards)
-                else:
-                    task_score = 0.01
-                task_score = max(0.01, min(0.99, task_score))
-                log_end(
-                    success=task_score >= 0.5,
-                    steps=task_step,
-                    score=task_score,
-                    rewards=task_rewards,
-                )
+            # Detect task transition
+            task_changed = (new_task != current_task) and (new_task != "unknown")
+
+            if task_changed or is_done:
+                # End current task
+                task_score = clamp(sum(task_rewards) / len(task_rewards)) if task_rewards else 0.01
+                log_end(task_score >= 0.5, task_step, task_score, task_rewards)
 
                 if is_done:
                     break
 
-                # Move to next task — emit new [START]
-                current_task_id = new_task_id
+                # Start next task
+                current_task = new_task
                 task_step = 0
                 task_rewards = []
                 task_history = []
-                log_start(task=current_task_id, env=benchmark, model=model_name)
+                log_start(task=current_task, env=benchmark, model=model_name)
+
+            result = step_result
 
     except Exception as e:
-        print(f"[DEBUG] Error during inference: {e}", flush=True)
-        # Emit a safe [END] so the validator has something to parse
-        log_end(success=False, steps=0, score=0.01, rewards=[0.01])
+        print(f"[DEBUG] Fatal error: {e}", flush=True)
+        log_end(False, 0, 0.01, [0.01])
     finally:
         try:
             await env.close()
@@ -250,4 +214,4 @@ if __name__ == "__main__":
     try:
         asyncio.run(main())
     except Exception:
-        print(f"[END] success=false steps=0 score=0.01 rewards=0.01", flush=True)
+        print("[END] success=false steps=0 score=0.01 rewards=0.01", flush=True)

server/config_debug_environment.py

@@ -56,25 +56,34 @@ class ConfigDebugEnvironment(Environment):
         task_id = self._current_task_id()
         task = get_task(task_id)
 
-        # Run the grader - now returns FLOAT ONLY for validator compatibility
+        # Run the grader - returns (reward, error_message, bugs_fixed) tuple
         grader_result = task.grader(action.fixed_config)
-        
-        # Handle result gracefully (should be float)
-        if isinstance(grader_result, float):
-            reward = grader_result
-        elif isinstance(grader_result, tuple) and len(grader_result) > 0:
-            # Fallback for legacy tuple format
-            reward = grader_result[0]
+
+        # Parse grader result
+        if isinstance(grader_result, tuple) and len(grader_result) >= 3:
+            reward = float(grader_result[0])
+            error_message = str(grader_result[1])
+            bugs_fixed = list(grader_result[2])
+        elif isinstance(grader_result, tuple) and len(grader_result) >= 1:
+            reward = float(grader_result[0])
+            error_message = ""
+            bugs_fixed = []
+        elif isinstance(grader_result, (int, float)):
+            reward = float(grader_result)
+            error_message = ""
+            bugs_fixed = []
         else:
-            reward = 0.01  # Safe default
-        
+            reward = 0.01
+            error_message = "Grader returned unexpected format"
+            bugs_fixed = []
+
         reward = max(0.01, min(0.99, reward))
 
         self.current_step += 1
         self._global_step += 1
-        self.bugs_found_so_far = 0  # Default since grader no longer returns this
+        self.bugs_found_so_far = len(bugs_fixed)
         self.previous_reward = round(reward, 4)
-        self.current_error_message = ""  # Default since grader no longer returns this
+        self.current_error_message = error_message
 
         # Check if task is complete
         task_done = reward >= 0.99 or self.current_step >= MAX_STEPS_PER_TASK
@@ -120,7 +129,6 @@ class ConfigDebugEnvironment(Environment):
             is_done=self._done,
             tasks_completed=list(self.tasks_completed),
             tasks_remaining=tasks_remaining,
-            # Enhanced RL signals
             bugs_found_so_far=self.bugs_found_so_far,
             current_error_message=self.current_error_message,
             progress_ratio=round(progress_ratio, 2),
@@ -153,4 +161,4 @@ class ConfigDebugEnvironment(Environment):
             previous_reward=self.previous_reward,
             done=self._done,
             reward=self.previous_reward,
-        )
+        )

server/tasks/task_registry.py

@@ -4,7 +4,7 @@ from server.tasks import task1_json, task2_yaml, task3_dockerfile
 from server.tasks import task4_compose, task5_k8s, task6_github_actions, task7_nginx
 
 # INTERNAL USE: Import directly from raw grader files (return tuples)
-# grader_api.py returns float-only (for validator/openenv.yaml)
+# grader_api.py is separate and returns float-only (for openenv.yaml validator)
 from server.graders.json_grader import grade_task1
 from server.graders.yaml_grader import grade_task2
 from server.graders.dockerfile_grader import grade_task3
@@ -15,7 +15,7 @@ from server.graders.nginx_grader import grade_task7
 
 
 def _clamp_grader(fn):
-    """Wrap raw grader to clamp reward to (0.01, 0.99)."""
+    """Wrap raw grader to clamp reward to (0.01, 0.99) while preserving tuple return."""
     def wrapper(submitted_config: str) -> Tuple[float, str, List[str]]:
         reward, error_msg, bugs_fixed = fn(submitted_config)
         reward = max(0.01, min(0.99, float(reward)))
@@ -64,4 +64,4 @@ def get_task(task_id: str) -> TaskInfo:
 
 
 def get_all_task_ids() -> List[str]:
-    return list(TASK_ORDER)
+    return list(TASK_ORDER)