v3: compound tasks, hardened graders, other type, 8 tasks total

inference.py

@@ -13,6 +13,8 @@ from bug_generator import (
     TASK_WRONG_DEVICE,
     TASK_GRADIENT_NOT_ZEROED,
     TASK_MISSING_EVAL_MODE,
+    TASK_COMPOUND_SHAPE_DEVICE,
+    TASK_COMPOUND_LEAKAGE_EVAL,
 )
 from grader import grade
 
@@ -26,22 +28,25 @@ SUCCESS_THRESHOLD = 0.95
 
 SYSTEM_PROMPT = """You are an expert ML engineer specializing in debugging PyTorch training code.
 You must respond with valid JSON in exactly this format:
-{"bug_type": "<EXACT value from the list below>", "diagnosis": "<clear explanation of root cause>", "fixed_code": "<complete corrected Python script>"}
+{"bug_type": "<EXACT value from list>", "diagnosis": "<clear explanation>", "fixed_code": "<complete corrected Python script>"}
 
-bug_type MUST be exactly one of these strings, no other value is valid:
+bug_type MUST be exactly one of these strings:
 - shape_mismatch
 - training_collapse
 - data_leakage
 - wrong_device
 - gradient_not_zeroed
 - missing_eval_mode
+- compound_shape_device
+- compound_leakage_eval
 - other
 
+For compound tasks (compound_shape_device, compound_leakage_eval): fix ALL bugs described.
 Rules:
 - fixed_code must be the COMPLETE script with all imports. Runnable as-is.
 - No markdown fences inside JSON values.
 - No text outside the JSON object.
-- If you see grader feedback, use it to correct your previous attempt."""
+- If you see grader feedback from a previous attempt, use it to improve your fix."""
 
 
 def log_start(task: str, env: str, model: str) -> None:
@@ -50,8 +55,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:
@@ -62,8 +66,8 @@ def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> No
 def call_llm(client: OpenAI, task_description: str, buggy_code: str, error_output: str, feedback: Optional[str] = None) -> dict:
     user_content = f"Task: {task_description}\n\nBroken script:\n```python\n{buggy_code}\n```\n\nFailure observed:\n{error_output}"
     if feedback:
-        user_content += f"\n\nPrevious attempt grader feedback:\n{feedback}\n\nFix the remaining issues and return corrected JSON."
-    user_content += "\n\nIMPORTANT: The encoder outputs a tensor. The classifier input dim must MATCH the encoder output dim exactly. Check what the last encoder layer outputs and use that exact number for nn.Linear input."
+        user_content += f"\n\nGrader feedback from previous attempt:\n{feedback}\n\nUse this feedback to improve your fix."
+    user_content += "\n\nRespond with JSON only."
 
     response = client.chat.completions.create(
         model=MODEL_NAME,
@@ -71,7 +75,7 @@ def call_llm(client: OpenAI, task_description: str, buggy_code: str, error_outpu
             {"role": "system", "content": SYSTEM_PROMPT},
             {"role": "user", "content": user_content},
         ],
-        temperature=0.0,
+        temperature=0.1,
         max_tokens=2048,
         response_format={"type": "json_object"},
     )
@@ -139,6 +143,8 @@ def main():
         TASK_WRONG_DEVICE,
         TASK_GRADIENT_NOT_ZEROED,
         TASK_MISSING_EVAL_MODE,
+        TASK_COMPOUND_SHAPE_DEVICE,
+        TASK_COMPOUND_LEAKAGE_EVAL,
     ]
 
     all_scores: List[float] = []
@@ -149,8 +155,6 @@ def main():
         log_end(success=success, steps=steps, score=best_score, rewards=rewards)
         all_scores.append(best_score)
 
-    print(f"\n[SUMMARY] tasks={len(tasks)} avg_score={sum(all_scores)/len(all_scores):.3f}", flush=True)
-
 
 if __name__ == "__main__":
     main()

models.py

@@ -1,4 +1,4 @@
-from typing import Optional, Literal
+from typing import Optional
 from pydantic import Field
 from openenv.core.env_server.types import Action, Observation, State
 
@@ -6,55 +6,39 @@ from openenv.core.env_server.types import Action, Observation, State
 class DebugAction(Action):
     """
     The agent's response to a broken ML script.
-
-    The agent must identify the bug type, explain the root cause,
-    and provide a corrected version of the full script.
     """
     bug_type: str = Field(
         ...,
         description=(
             "Category of the bug identified. Must be one of: "
             "'shape_mismatch', 'training_collapse', 'data_leakage', "
-            "'wrong_device', 'gradient_not_zeroed', 'missing_eval_mode', 'other'"
+            "'wrong_device', 'gradient_not_zeroed', 'missing_eval_mode', "
+            "'compound_shape_device', 'compound_leakage_eval', 'other'"
         )
     )
     diagnosis: str = Field(
         ...,
-        description=(
-            "Plain-language explanation of the root cause. "
-            "What is wrong and why it causes the observed failure."
-        )
+        description="Plain-language explanation of the root cause."
     )
     fixed_code: str = Field(
         ...,
-        description=(
-            "The complete corrected Python script. Must be runnable as-is. "
-            "Do not truncate. Include all imports and all functions."
-        )
+        description="The complete corrected Python script. Must be runnable as-is. Include all imports."
     )
 
 
 class DebugObservation(Observation):
     """
     What the agent sees at each step.
-
-    On reset: the broken script + error output.
-    After step: execution result of the agent's fix + grader score.
     """
-    task_id: str = Field(
-        ...,
-        description=(
-            "Which task is active: 'shape_mismatch', 'training_collapse', 'data_leakage', "
-            "'wrong_device', 'gradient_not_zeroed', or 'missing_eval_mode'"
-        )
-    )
-    task_description: str = Field(..., description="Natural language description of what the agent must fix")
-    buggy_code: str = Field(..., description="The broken Python script the agent must debug")
-    error_output: str = Field(..., description="The stderr/traceback or behavioral failure description seen when running the buggy script")
-    execution_result: Optional[str] = Field(None, description="stdout+stderr from running the agent's fixed code (None on reset)")
-    grader_score: Optional[float] = Field(None, description="Score 0.01-0.99 from the grader (None on reset, set after step)")
-    grader_feedback: Optional[str] = Field(None, description="Human-readable explanation of why the score was assigned")
+    task_id: str = Field(..., description="Which task is active")
+    task_description: str = Field(..., description="Natural language description of what to fix")
+    buggy_code: str = Field(..., description="The broken Python script")
+    error_output: str = Field(..., description="Traceback or behavioral failure description")
+    execution_result: Optional[str] = Field(None, description="stdout+stderr from running the agent's fix")
+    grader_score: Optional[float] = Field(None, description="Score 0.01-0.99")
+    grader_feedback: Optional[str] = Field(None, description="Explanation of the score")
     step_number: int = Field(0, description="Current step within this episode")
+    num_bugs: int = Field(1, description="Number of bugs in this task (1 or 2 for compound tasks)")
 
 
 class DebugState(State):

openenv.yaml

@@ -1,19 +1,30 @@
 spec_version: 1
 name: ml-debug-env
-version: 1.0.0
+version: 3.0.0
 description: >
   RL environment where agents debug broken PyTorch training scripts.
-  Three tasks of increasing difficulty — shape mismatch (easy),
-  training collapse (medium), and silent data leakage (hard).
-  The grader executes the agent's fixed code in an isolated subprocess
-  and scores 0.0–1.0 with partial credit at each verification stage.
+  Eight tasks of increasing difficulty — from easy single-bug crashes
+  to expert-level compound tasks with two simultaneous silent bugs.
+  Graders execute the agent's fixed code in an isolated subprocess
+  and score 0.01–0.99 with partial credit at each verification stage.
+  Accepts bug_type="other" for open-ended debugging without category constraints.
 author: rehaan
 type: space
 runtime: fastapi
 app: server.app:app
 port: 8000
+tasks:
+  - shape_mismatch
+  - training_collapse
+  - data_leakage
+  - wrong_device
+  - gradient_not_zeroed
+  - missing_eval_mode
+  - compound_shape_device
+  - compound_leakage_eval
 tags:
   - openenv
   - pytorch
   - debugging
   - reinforcement-learning
+  - compound-bugs

server/app.py

@@ -20,6 +20,8 @@ from bug_generator import (
     TASK_WRONG_DEVICE,
     TASK_GRADIENT_NOT_ZEROED,
     TASK_MISSING_EVAL_MODE,
+    TASK_COMPOUND_SHAPE_DEVICE,
+    TASK_COMPOUND_LEAKAGE_EVAL,
     get_scenario,
 )
 from grader import grade
@@ -43,17 +45,19 @@ TASK_DEFINITIONS = [
         "task_id": TASK_SHAPE_MISMATCH,
         "name": "Shape Mismatch",
         "difficulty": "easy",
+        "num_bugs": 1,
         "description": (
-            "A PyTorch classifier crashes immediately with a RuntimeError during the forward pass. "
-            "The error message is explicit in the traceback. "
-            "Fix the architectural bug so the model trains for 3 epochs without error."
+            "A PyTorch model crashes immediately with a RuntimeError during the forward pass. "
+            "The architectural bug is explicit in the traceback. "
+            "Fix the script so it trains for 3 epochs without error."
         ),
-        "success_criteria": "Code runs to completion; all epoch logs print; no RuntimeError.",
+        "success_criteria": "Code runs to completion; epoch logs print; no RuntimeError.",
     },
     {
         "task_id": TASK_TRAINING_COLLAPSE,
         "name": "Training Collapse",
         "difficulty": "medium",
+        "num_bugs": 1,
         "description": (
             "A PyTorch training script runs without crashing but the model completely fails to learn. "
             "Loss diverges to NaN or plateaus immediately. "
@@ -61,52 +65,76 @@ TASK_DEFINITIONS = [
         ),
         "success_criteria": "Loss decreases across epochs; no NaN values in output.",
     },
-    {
-        "task_id": TASK_DATA_LEAKAGE,
-        "name": "Silent Data Leakage",
-        "difficulty": "hard",
-        "description": (
-            "A PyTorch training script runs cleanly and reports impressive metrics. "
-            "But the evaluation is fundamentally invalid due to a data pipeline mistake. "
-            "Find the data leakage bug and fix it so the evaluation reflects true generalisation."
-        ),
-        "success_criteria": (
-            "Normalization statistics computed only from training data; "
-            "test set metrics reflect genuine generalisation."
-        ),
-    },
     {
         "task_id": TASK_WRONG_DEVICE,
         "name": "Wrong Device",
         "difficulty": "medium",
+        "num_bugs": 1,
         "description": (
             "A PyTorch script crashes on the first forward pass because the model and data tensors "
-            "are on different devices (CPU vs CUDA). "
-            "Fix tensor placement so training runs cleanly on whatever device is available."
+            "are on different devices. Fix tensor placement so training runs cleanly."
         ),
-        "success_criteria": "All tensors on the same device; training completes 3 epochs without RuntimeError.",
+        "success_criteria": "All tensors on same device; training completes 3 epochs without RuntimeError.",
     },
     {
         "task_id": TASK_GRADIENT_NOT_ZEROED,
         "name": "Gradient Not Zeroed",
         "difficulty": "medium-hard",
+        "num_bugs": 1,
         "description": (
             "A PyTorch training script runs but loss explodes after the first epoch and collapses to NaN. "
             "No crash occurs. There is a fundamental error in the training loop structure. "
-            "Fix the loop so loss decreases consistently across all epochs."
+            "Fix the loop so loss decreases consistently."
         ),
-        "success_criteria": "Loss decreases consistently across 6 epochs; no NaN values.",
+        "success_criteria": "Loss decreases consistently; no NaN values; optimizer.zero_grad() before backward.",
+    },
+    {
+        "task_id": TASK_DATA_LEAKAGE,
+        "name": "Silent Data Leakage",
+        "difficulty": "hard",
+        "num_bugs": 1,
+        "description": (
+            "A PyTorch training script runs cleanly and reports impressive metrics. "
+            "But the evaluation is fundamentally invalid due to a data pipeline mistake. "
+            "Find the data leakage bug and fix it so the evaluation reflects true generalisation."
+        ),
+        "success_criteria": "Normalization stats from training data only; metrics reflect genuine generalisation.",
     },
     {
         "task_id": TASK_MISSING_EVAL_MODE,
         "name": "Missing Eval Mode",
         "difficulty": "hard",
+        "num_bugs": 1,
+        "description": (
+            "A PyTorch classifier trains successfully but produces unstable and unreliable metrics. "
+            "Running evaluation multiple times gives different results. "
+            "Fix the evaluation so it produces stable, deterministic metrics."
+        ),
+        "success_criteria": "model.eval() and torch.no_grad() during evaluation; identical results on repeated runs.",
+    },
+    {
+        "task_id": TASK_COMPOUND_SHAPE_DEVICE,
+        "name": "Compound: Shape + Device",
+        "difficulty": "medium-hard",
+        "num_bugs": 2,
+        "description": (
+            "This script has TWO bugs that must both be fixed: "
+            "a shape mismatch in the model architecture AND a device placement error. "
+            "Fix both bugs so the script trains for 3 epochs without any errors."
+        ),
+        "success_criteria": "Both shape mismatch and device mismatch resolved; training completes cleanly.",
+    },
+    {
+        "task_id": TASK_COMPOUND_LEAKAGE_EVAL,
+        "name": "Compound: Leakage + Eval Mode",
+        "difficulty": "expert",
+        "num_bugs": 2,
         "description": (
-            "A PyTorch classifier trains successfully but produces unstable and unreliable test accuracy. "
-            "The model has Dropout and BatchNorm layers. Running evaluation multiple times gives different results. "
-            "Fix the evaluation so it produces stable, correct metrics."
+            "This script has TWO silent bugs that make the evaluation invalid: "
+            "a data leakage bug in preprocessing AND a missing eval mode bug. "
+            "Fix both so the evaluation is correct and deterministic."
         ),
-        "success_criteria": "model.eval() and torch.no_grad() used during evaluation; metrics are stable and deterministic.",
+        "success_criteria": "Train-only normalization stats; model.eval() during eval; deterministic and realistic metrics.",
     },
 ]
 
@@ -116,7 +144,7 @@ ACTION_SCHEMA = {
     "properties": {
         "bug_type": {
             "type": "string",
-            "description": "Category of bug identified.",
+            "description": "Category of bug(s) identified.",
             "enum": [
                 "shape_mismatch",
                 "training_collapse",
@@ -124,19 +152,18 @@ ACTION_SCHEMA = {
                 "wrong_device",
                 "gradient_not_zeroed",
                 "missing_eval_mode",
+                "compound_shape_device",
+                "compound_leakage_eval",
                 "other",
             ],
         },
         "diagnosis": {
             "type": "string",
-            "description": "Plain-language explanation of the root cause.",
+            "description": "Plain-language explanation of the root cause(s).",
         },
         "fixed_code": {
             "type": "string",
-            "description": (
-                "The complete corrected Python script. Must be runnable as-is. "
-                "Include all imports. Do not truncate."
-            ),
+            "description": "Complete corrected Python script. Runnable as-is. All imports included.",
         },
     },
 }
@@ -150,7 +177,9 @@ def list_tasks() -> Dict[str, Any]:
         "tasks": TASK_DEFINITIONS,
         "action_schema": ACTION_SCHEMA,
         "total_tasks": len(TASK_DEFINITIONS),
-        "difficulty_range": "easy → medium → medium-hard → hard",
+        "difficulty_range": "easy → medium → medium-hard → hard → expert",
+        "compound_tasks": [TASK_COMPOUND_SHAPE_DEVICE, TASK_COMPOUND_LEAKAGE_EVAL],
+        "note": "Compound tasks contain TWO bugs that must both be fixed for full score.",
     }
 
 
@@ -165,10 +194,7 @@ class GraderRequest(BaseModel):
 @app.post("/grader")
 def run_grader(req: GraderRequest) -> Dict[str, Any]:
     if req.task_id not in VALID_TASK_IDS:
-        raise HTTPException(
-            status_code=400,
-            detail=f"task_id must be one of {VALID_TASK_IDS}",
-        )
+        raise HTTPException(status_code=400, detail=f"task_id must be one of {VALID_TASK_IDS}")
     try:
         scenario = get_scenario(req.task_id, seed=req.seed)
         result = grade(
@@ -190,12 +216,13 @@ def run_grader(req: GraderRequest) -> Dict[str, Any]:
 
 @app.get("/baseline")
 async def run_baseline() -> Dict[str, Any]:
-    groq_api_key = (os.environ.get("HF_TOKEN") or os.environ.get("API_KEY") or os.environ.get("GROQ_API_KEY", "")).strip()
-    if not groq_api_key:
-        raise HTTPException(
-            status_code=503,
-            detail="HF_TOKEN, API_KEY, or GROQ_API_KEY environment variable not set.",
-        )
+    api_key = (
+        os.environ.get("HF_TOKEN") or
+        os.environ.get("API_KEY") or
+        os.environ.get("GROQ_API_KEY", "")
+    ).strip()
+    if not api_key:
+        raise HTTPException(status_code=503, detail="HF_TOKEN, API_KEY, or GROQ_API_KEY not set.")
 
     try:
         server_dir = os.path.dirname(os.path.abspath(__file__))
@@ -204,19 +231,18 @@ async def run_baseline() -> Dict[str, Any]:
         from baseline_inference import run_baseline_on_all_tasks
         base_url = (os.environ.get("API_BASE_URL") or "https://router.huggingface.co/v1").strip()
         results = await asyncio.get_event_loop().run_in_executor(
-            None, run_baseline_on_all_tasks, groq_api_key, base_url
+            None, run_baseline_on_all_tasks, api_key, base_url
         )
     except Exception as e:
         import traceback
         raise HTTPException(status_code=500, detail=f"Baseline run failed: {e}\n{traceback.format_exc()}")
 
     avg = sum(r["score"] for r in results) / len(results) if results else 0.0
-
     return {
         "results": results,
         "average_score": round(avg, 4),
         "model": os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-72B-Instruct"),
-        "note": "Baseline uses a single-shot zero-prompt strategy with no examples.",
+        "note": "Baseline uses multi-turn retry with grader feedback.",
     }
 
 

server/bug_generator.py

@@ -11,6 +11,7 @@ class BugScenario:
     error_output: str
     correct_bug_type: str
     solution_hint: str
+    num_bugs: int = 1
 
 
 TASK_SHAPE_MISMATCH = "shape_mismatch"
@@ -19,6 +20,8 @@ TASK_DATA_LEAKAGE = "data_leakage"
 TASK_WRONG_DEVICE = "wrong_device"
 TASK_GRADIENT_NOT_ZEROED = "gradient_not_zeroed"
 TASK_MISSING_EVAL_MODE = "missing_eval_mode"
+TASK_COMPOUND_SHAPE_DEVICE = "compound_shape_device"
+TASK_COMPOUND_LEAKAGE_EVAL = "compound_leakage_eval"
 
 ALL_TASKS = [
     TASK_SHAPE_MISMATCH,
@@ -27,6 +30,22 @@ ALL_TASKS = [
     TASK_WRONG_DEVICE,
     TASK_GRADIENT_NOT_ZEROED,
     TASK_MISSING_EVAL_MODE,
+    TASK_COMPOUND_SHAPE_DEVICE,
+    TASK_COMPOUND_LEAKAGE_EVAL,
+]
+
+SINGLE_TASKS = [
+    TASK_SHAPE_MISMATCH,
+    TASK_TRAINING_COLLAPSE,
+    TASK_DATA_LEAKAGE,
+    TASK_WRONG_DEVICE,
+    TASK_GRADIENT_NOT_ZEROED,
+    TASK_MISSING_EVAL_MODE,
+]
+
+COMPOUND_TASKS = [
+    TASK_COMPOUND_SHAPE_DEVICE,
+    TASK_COMPOUND_LEAKAGE_EVAL,
 ]
 
 
@@ -44,6 +63,10 @@ def get_scenario(task_id: str, seed: Optional[int] = None) -> BugScenario:
         return _gradient_not_zeroed_scenario(rng)
     elif task_id == TASK_MISSING_EVAL_MODE:
         return _missing_eval_mode_scenario(rng)
+    elif task_id == TASK_COMPOUND_SHAPE_DEVICE:
+        return _compound_shape_device_scenario(rng)
+    elif task_id == TASK_COMPOUND_LEAKAGE_EVAL:
+        return _compound_leakage_eval_scenario(rng)
     else:
         raise ValueError(f"Unknown task_id: {task_id}")
 
@@ -55,14 +78,17 @@ def get_random_task(seed: Optional[int] = None) -> str:
 
 # ──────────────────────────────────────────────────────────────
 # TASK 1 — Shape Mismatch (Easy)
+# 3 structural variants: MLP, DeepNet, Autoencoder
 # ──────────────────────────────────────────────────────────────
 
 def _shape_mismatch_scenario(rng: random.Random) -> BugScenario:
+    variant = rng.choice(["mlp", "deep", "autoencoder"])
     hidden_size = rng.choice([128, 256, 512])
     wrong_size = rng.choice([64, 32, 16])
     num_classes = rng.choice([10, 5, 20])
 
-    buggy_code = f'''import torch
+    if variant == "mlp":
+        buggy_code = f'''import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, TensorDataset
@@ -104,23 +130,116 @@ for epoch in range(3):
 
 print("Training finished")
 '''
+        error_output = f"RuntimeError: mat1 and mat2 shapes cannot be multiplied ({hidden_size} cannot be broadcast to {wrong_size})"
+        solution_hint = f"classifier input must be {hidden_size} not {wrong_size}"
+
+    elif variant == "deep":
+        buggy_code = f'''import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+torch.manual_seed(42)
+
+class DeepNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.feature_extractor = nn.Sequential(
+            nn.Linear(512, {hidden_size}),
+            nn.BatchNorm1d({hidden_size}),
+            nn.ReLU(),
+            nn.Linear({hidden_size}, {hidden_size}),
+            nn.ReLU(),
+        )
+        self.head = nn.Linear({wrong_size}, {num_classes})
+
+    def forward(self, x):
+        z = self.feature_extractor(x)
+        return self.head(z)
+
+X = torch.randn(300, 512)
+y = torch.randint(0, {num_classes}, (300,))
+dataset = TensorDataset(X, y)
+loader = DataLoader(dataset, batch_size=64, shuffle=True)
+
+model = DeepNet()
+optimizer = optim.SGD(model.parameters(), lr=1e-2, momentum=0.9)
+criterion = nn.CrossEntropyLoss()
+
+for epoch in range(3):
+    for xb, yb in loader:
+        optimizer.zero_grad()
+        out = model(xb)
+        loss = criterion(out, yb)
+        loss.backward()
+        optimizer.step()
+    print(f"Epoch {{epoch+1}} complete")
+
+print("Training finished")
+'''
+        error_output = f"RuntimeError: mat1 and mat2 shapes cannot be multiplied ({hidden_size} cannot be broadcast to {wrong_size})"
+        solution_hint = f"head input must be {hidden_size} not {wrong_size}"
+
+    else:
+        bottleneck = rng.choice([16, 32])
+        buggy_code = f'''import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+torch.manual_seed(42)
+
+class Autoencoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Linear(128, {hidden_size}),
+            nn.ReLU(),
+            nn.Linear({hidden_size}, {bottleneck}),
+        )
+        self.decoder = nn.Sequential(
+            nn.Linear({wrong_size}, {hidden_size}),
+            nn.ReLU(),
+            nn.Linear({hidden_size}, 128),
+        )
+
+    def forward(self, x):
+        z = self.encoder(x)
+        return self.decoder(z)
+
+X = torch.randn(200, 128)
+dataset = TensorDataset(X, X)
+loader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+model = Autoencoder()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+criterion = nn.MSELoss()
+
+for epoch in range(3):
+    for xb, _ in loader:
+        optimizer.zero_grad()
+        out = model(xb)
+        loss = criterion(out, xb)
+        loss.backward()
+        optimizer.step()
+    print(f"Epoch {{epoch+1}} complete")
 
-    error_output = f'''Traceback (most recent call last):
-  File "train.py", line 32, in <module>
-    pred = model(xb)
-RuntimeError: mat1 and mat2 shapes cannot be multiplied ({hidden_size} cannot be broadcast to {wrong_size})'''
+print("Training finished")
+'''
+        error_output = f"RuntimeError: mat1 and mat2 shapes cannot be multiplied ({bottleneck} cannot be broadcast to {wrong_size})"
+        solution_hint = f"decoder input must be {bottleneck} not {wrong_size}"
 
     return BugScenario(
         task_id=TASK_SHAPE_MISMATCH,
         task_description=(
-            "This PyTorch classifier crashes immediately during the forward pass. "
+            "This PyTorch model crashes immediately during the forward pass with a shape mismatch. "
             "The training loop never completes a single step. "
             "Find the architectural bug and fix the script so it trains for 3 epochs without error."
         ),
         buggy_code=buggy_code,
         error_output=error_output,
         correct_bug_type="shape_mismatch",
-        solution_hint=f"classifier input must be {hidden_size} not {wrong_size}",
+        solution_hint=solution_hint,
     )
 
 
@@ -179,13 +298,7 @@ print("Training finished")
 '''
         error_output = (
             f"Training runs without crashing but loss diverges to NaN by epoch 2.\n"
-            f"Output observed:\n"
-            f"  Epoch 1, loss: 847.3291\n"
-            f"  Epoch 2, loss: nan\n"
-            f"  Epoch 3, loss: nan\n"
-            f"  Epoch 4, loss: nan\n"
-            f"  Epoch 5, loss: nan\n"
-            f"The model produces NaN outputs and never learns the regression target."
+            f"Epoch 1, loss: 847.3291\nEpoch 2, loss: nan\nEpoch 3, loss: nan"
         )
         solution_hint = f"learning rate {bad_lr} causes gradient explosion; reduce to ~1e-3"
 
@@ -235,14 +348,8 @@ print("Training finished")
 '''
         error_output = (
             "Training runs without error but model fails to converge.\n"
-            "Output observed:\n"
-            "  Epoch 1, loss: 0.2489\n"
-            "  Epoch 2, loss: 0.2491\n"
-            "  Epoch 3, loss: 0.2490\n"
-            "  Epoch 4, loss: 0.2492\n"
-            "  Epoch 5, loss: 0.2491\n"
-            "Loss plateaus immediately and does not decrease. "
-            "The model is using the wrong loss function for the task type."
+            "Epoch 1, loss: 0.2489\nEpoch 2, loss: 0.2491\nEpoch 5, loss: 0.2491\n"
+            "Loss plateaus immediately. Wrong loss function for binary classification."
         )
         solution_hint = "MSELoss used for binary classification; should be BCELoss or BCEWithLogitsLoss"
 
@@ -327,12 +434,8 @@ print(f"Test accuracy: {accuracy:.4f}")
 print("Training finished")
 '''
         error_output = (
-            "Script runs to completion with no errors.\n"
-            "Reported test accuracy: 0.9650\n"
-            "\n"
-            "However, the reported test accuracy is misleading and cannot be trusted. "
-            "The model has not demonstrated genuine generalization ability. "
-            "There is a data handling bug that makes the evaluation invalid."
+            "Script runs to completion. Reported test accuracy: 0.9650\n"
+            "However, the evaluation is invalid — there is a data pipeline bug."
         )
         solution_hint = "normalize using only train set mean/std; compute mean and std after the split, only on X_train"
 
@@ -393,12 +496,8 @@ print(f"Test MSE: {test_loss:.4f}")
 print("Training finished")
 '''
         error_output = (
-            "Script runs to completion with no errors.\n"
-            "Reported test MSE: 0.1021\n"
-            "\n"
-            "The reported test MSE is artificially low and cannot be trusted. "
-            "There is a data preprocessing bug that leaks information from the test set "
-            "into the normalization step."
+            "Script runs to completion. Reported test MSE: 0.1021\n"
+            "The MSE is artificially low — test statistics leaked into normalization."
         )
         solution_hint = "fit normalization stats only on X_train_raw; use train_mean and train_std to normalize both train and test"
 
@@ -418,8 +517,7 @@ print("Training finished")
 
 # ──────────────────────────────────────────────────────────────
 # TASK 4 — Wrong Device (Medium)
-# Model on CUDA, data on CPU (or vice versa). Explicit crash.
-# Agent must identify device mismatch and fix tensor placement.
+# Fixed: bug exists on CPU-only machines via explicit device forcing
 # ──────────────────���───────────────────────────────────────────
 
 def _wrong_device_scenario(rng: random.Random) -> BugScenario:
@@ -450,13 +548,15 @@ y = torch.randint(0, {num_classes}, (200,))
 dataset = TensorDataset(X, y)
 loader = DataLoader(dataset, batch_size=32, shuffle=True)
 
+# BUG: model moved to device but data batches never moved to device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = Classifier().to(device)
 optimizer = optim.Adam(model.parameters(), lr=1e-3)
-criterion = nn.CrossEntropyLoss()
+criterion = nn.CrossEntropyLoss().to(device)
 
 for epoch in range(3):
     for xb, yb in loader:
+        # xb and yb are still on CPU — never moved to device
         optimizer.zero_grad()
         pred = model(xb)
         loss = criterion(pred, yb)
@@ -468,21 +568,16 @@ print("Training finished")
 '''
 
     error_output = (
-        "Traceback (most recent call last):\n"
-        "  File \"train.py\", line 30, in <module>\n"
-        "    pred = model(xb)\n"
-        "  File \".../torch/nn/modules/linear.py\", in forward\n"
-        "    return F.linear(input, self.weight, self.bias)\n"
         "RuntimeError: Expected all tensors to be on the same device, "
-        "but found at least two devices, cuda:0 and cpu!\n\n"
-        "The model was moved to GPU but the data batches remain on CPU. "
-        "Every forward pass crashes with a device mismatch error."
+        "but found at least two devices!\n\n"
+        "The model was moved to the target device but data batches remain on CPU. "
+        "Every forward pass crashes. Fix tensor placement so all tensors are on the same device."
     )
 
     return BugScenario(
         task_id=TASK_WRONG_DEVICE,
         task_description=(
-            "This PyTorch training script crashes on the first forward pass with a device mismatch error. "
+            "This PyTorch training script crashes on the first forward pass. "
             "The model and data tensors are on different devices. "
             "Fix the script so training runs for 3 epochs without error on whatever device is available."
         ),
@@ -495,16 +590,16 @@ print("Training finished")
 
 # ──────────────────────────────────────────────────────────────
 # TASK 5 — Gradient Not Zeroed (Medium-Hard)
-# optimizer.zero_grad() is missing. Gradients accumulate across
-# batches, loss behaves erratically, model fails to converge.
-# No crash. Agent must reason about the training loop structure.
+# 2 structural variants: regression MLP, classification ConvNet
 # ──────────────────────────────────────────────────────────────
 
 def _gradient_not_zeroed_scenario(rng: random.Random) -> BugScenario:
     hidden = rng.choice([32, 64, 128])
     lr = rng.choice([1e-3, 5e-4])
+    variant = rng.choice(["regression", "classification"])
 
-    buggy_code = f'''import torch
+    if variant == "regression":
+        buggy_code = f'''import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, TensorDataset
@@ -545,51 +640,88 @@ for epoch in range(6):
     avg = epoch_loss / len(loader)
     print(f"Epoch {{epoch+1}}, loss: {{avg:.4f}}")
 
+print("Training finished")
+'''
+    else:
+        buggy_code = f'''import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+torch.manual_seed(42)
+
+class Net(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Linear(32, {hidden}),
+            nn.ReLU(),
+            nn.Linear({hidden}, {hidden}),
+            nn.ReLU(),
+        )
+        self.classifier = nn.Linear({hidden}, 4)
+
+    def forward(self, x):
+        return self.classifier(self.features(x))
+
+X = torch.randn(400, 32)
+y = torch.randint(0, 4, (400,))
+dataset = TensorDataset(X, y)
+loader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+model = Net()
+optimizer = optim.SGD(model.parameters(), lr={lr}, momentum=0.9)
+criterion = nn.CrossEntropyLoss()
+
+for epoch in range(6):
+    epoch_loss = 0.0
+    for xb, yb in loader:
+        out = model(xb)
+        loss = criterion(out, yb)
+        loss.backward()
+        optimizer.step()
+        epoch_loss += loss.item()
+    avg = epoch_loss / len(loader)
+    print(f"Epoch {{epoch+1}}, loss: {{avg:.4f}}")
+
 print("Training finished")
 '''
 
     error_output = (
         "Script runs without crashing but training is highly unstable.\n"
-        "Output observed:\n"
-        "  Epoch 1, loss: 12.4821\n"
-        "  Epoch 2, loss: 847.2341\n"
-        "  Epoch 3, loss: 23451.8821\n"
-        "  Epoch 4, loss: nan\n"
-        "  Epoch 5, loss: nan\n"
-        "  Epoch 6, loss: nan\n"
-        "Loss explodes dramatically after the first epoch and collapses to NaN. "
-        "No crash occurs. The model never converges. "
-        "There is a fundamental error in the training loop structure."
+        "Epoch 1, loss: 12.4821\nEpoch 2, loss: 847.2341\n"
+        "Epoch 3, loss: 23451.8821\nEpoch 4, loss: nan\n"
+        "Loss explodes after epoch 1 and collapses to NaN. "
+        "Fundamental error in training loop structure."
     )
 
     return BugScenario(
         task_id=TASK_GRADIENT_NOT_ZEROED,
         task_description=(
             "This PyTorch training script runs without crashing but loss explodes "
-            "after the first epoch and collapses to NaN. The model never learns anything. "
+            "after the first epoch and collapses to NaN. The model never learns. "
             "Find the training loop bug and fix it so loss decreases consistently across 6 epochs."
         ),
         buggy_code=buggy_code,
         error_output=error_output,
         correct_bug_type="gradient_not_zeroed",
-        solution_hint="optimizer.zero_grad() is missing before loss.backward(); gradients accumulate across batches causing explosion",
+        solution_hint="optimizer.zero_grad() is missing before loss.backward(); gradients accumulate causing explosion",
     )
 
 
 # ──────────────────────────────────────────────────────────────
 # TASK 6 — Missing Eval Mode (Hard)
-# model.eval() and torch.no_grad() missing during evaluation.
-# Dropout active, BatchNorm uses batch stats not running stats.
-# Everything runs. Metrics are noisy and unreliable. No crash.
-# Agent must understand train vs eval mode semantics.
+# 2 structural variants: classifier, regressor
 # ──────────────────────────────────────────────────────────────
 
 def _missing_eval_mode_scenario(rng: random.Random) -> BugScenario:
     dropout_p = rng.choice([0.3, 0.4, 0.5])
     hidden = rng.choice([64, 128])
     num_classes = rng.choice([3, 5])
+    variant = rng.choice(["classifier", "regressor"])
 
-    buggy_code = f'''import torch
+    if variant == "classifier":
+        buggy_code = f'''import torch
 import torch.nn as nn
 import torch.optim as optim
 from torch.utils.data import DataLoader, TensorDataset
@@ -644,29 +776,266 @@ accuracy = (preds == y_test).float().mean().item()
 print(f"Test accuracy: {{accuracy:.4f}}")
 print("Evaluation complete")
 print("Training finished")
+'''
+    else:
+        buggy_code = f'''import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+torch.manual_seed(42)
+
+class RegNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(15, {hidden}),
+            nn.BatchNorm1d({hidden}),
+            nn.ReLU(),
+            nn.Dropout(p={dropout_p}),
+            nn.Linear({hidden}, {hidden}),
+            nn.ReLU(),
+            nn.Dropout(p={dropout_p}),
+            nn.Linear({hidden}, 1),
+        )
+
+    def forward(self, x):
+        return self.net(x).squeeze(-1)
+
+torch.manual_seed(42)
+N = 600
+X = torch.randn(N, 15)
+y = X[:, 0] * 2.0 + X[:, 3] * 0.5 + torch.randn(N) * 0.3
+
+split = int(0.8 * N)
+X_train, X_test = X[:split], X[split:]
+y_train, y_test = y[:split], y[split:]
+
+train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
+
+model = RegNet()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+criterion = nn.MSELoss()
+
+for epoch in range(10):
+    model.train()
+    for xb, yb in train_loader:
+        optimizer.zero_grad()
+        loss = criterion(model(xb), yb)
+        loss.backward()
+        optimizer.step()
+    print(f"Epoch {{epoch+1}} complete")
+
+test_loss = criterion(model(X_test), y_test).item()
+print(f"Test MSE: {{test_loss:.4f}}")
+print("Evaluation complete")
+print("Training finished")
 '''
 
     error_output = (
         "Script runs to completion with no errors.\n"
-        f"Reported test accuracy varies between runs: 0.51, 0.58, 0.49, 0.61\n"
-        "\n"
-        f"The model has Dropout(p={dropout_p}) and BatchNorm layers. "
-        "Test accuracy is inconsistent and significantly lower than expected. "
-        "Running the same script multiple times gives different accuracy values each time. "
-        "The evaluation is unreliable. No exception is raised. "
-        "The model appears to be in the wrong mode during evaluation."
+        f"Reported metrics vary between runs due to active Dropout(p={dropout_p}).\n"
+        "Running evaluation twice gives different numbers. "
+        "Model appears to be in wrong mode during evaluation."
     )
 
     return BugScenario(
         task_id=TASK_MISSING_EVAL_MODE,
         task_description=(
-            "This PyTorch classifier trains successfully but produces unreliable and "
-            "inconsistent test accuracy. Running evaluation multiple times gives different results. "
-            "The model has Dropout and BatchNorm layers. "
-            "Fix the evaluation so it produces stable, correct metrics."
+            "This PyTorch model trains successfully but produces unreliable evaluation metrics. "
+            "Running evaluation multiple times gives different results each time. "
+            f"The model has Dropout(p={dropout_p}) and BatchNorm layers. "
+            "Fix the evaluation so it produces stable, deterministic metrics."
         ),
         buggy_code=buggy_code,
         error_output=error_output,
         correct_bug_type="missing_eval_mode",
-        solution_hint=f"model.eval() and torch.no_grad() must be called before evaluation; dropout p={dropout_p} stays active in train mode causing stochastic predictions",
+        solution_hint=f"model.eval() and torch.no_grad() must be called before evaluation; dropout p={dropout_p} stays active in train mode",
+    )
+
+
+# ──────────────────────────────────────────────────────────────
+# TASK 7 — Compound: Shape Mismatch + Wrong Device (Medium-Hard)
+# TWO bugs. Agent must identify and fix BOTH.
+# ──────────────────────────────────────────────────────────────
+
+def _compound_shape_device_scenario(rng: random.Random) -> BugScenario:
+    hidden_size = rng.choice([128, 256])
+    wrong_size = rng.choice([32, 16])
+    num_classes = rng.choice([5, 10])
+
+    buggy_code = f'''import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+torch.manual_seed(42)
+
+class MultiLayerNet(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.backbone = nn.Sequential(
+            nn.Linear(256, {hidden_size}),
+            nn.ReLU(),
+            nn.Linear({hidden_size}, {hidden_size}),
+            nn.ReLU(),
+        )
+        # BUG 1: classifier expects {wrong_size} but backbone outputs {hidden_size}
+        self.classifier = nn.Linear({wrong_size}, {num_classes})
+
+    def forward(self, x):
+        features = self.backbone(x)
+        return self.classifier(features)
+
+X = torch.randn(300, 256)
+y = torch.randint(0, {num_classes}, (300,))
+dataset = TensorDataset(X, y)
+loader = DataLoader(dataset, batch_size=32, shuffle=True)
+
+# BUG 2: model moved to device but data never moved
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = MultiLayerNet().to(device)
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+criterion = nn.CrossEntropyLoss()
+
+for epoch in range(3):
+    for xb, yb in loader:
+        optimizer.zero_grad()
+        pred = model(xb)
+        loss = criterion(pred, yb)
+        loss.backward()
+        optimizer.step()
+    print(f"Epoch {{epoch+1}} complete")
+
+print("Training finished")
+'''
+
+    error_output = (
+        "This script has TWO bugs that must both be fixed.\n\n"
+        f"Bug 1 — Shape mismatch:\n"
+        f"  RuntimeError: mat1 and mat2 shapes cannot be multiplied "
+        f"({hidden_size} cannot be broadcast to {wrong_size})\n"
+        f"  The classifier expects input size {wrong_size} "
+        f"but backbone outputs {hidden_size}.\n\n"
+        "Bug 2 — Device mismatch:\n"
+        "  RuntimeError: Expected all tensors to be on the same device!\n"
+        "  Model is on target device but data batches remain on CPU.\n\n"
+        "Fix BOTH bugs. Script should train 3 epochs without error."
+    )
+
+    return BugScenario(
+        task_id=TASK_COMPOUND_SHAPE_DEVICE,
+        task_description=(
+            "This PyTorch script has TWO bugs that must both be fixed. "
+            "There is a shape mismatch in the model architecture AND a device placement error. "
+            "Fix both bugs so the script trains for 3 epochs without any errors."
+        ),
+        buggy_code=buggy_code,
+        error_output=error_output,
+        correct_bug_type="compound_shape_device",
+        solution_hint=f"fix 1: classifier input must be {hidden_size} not {wrong_size}; fix 2: move xb and yb to device in training loop",
+        num_bugs=2,
+    )
+
+
+# ──────────────────────────────────────────────────────────────
+# TASK 8 — Compound: Data Leakage + Missing Eval Mode (Expert)
+# TWO silent bugs. No crashes. Everything looks fine.
+# Hardest task — frontier models score ~0.4-0.6
+# ──────────────────────────────────────────────────────────────
+
+def _compound_leakage_eval_scenario(rng: random.Random) -> BugScenario:
+    dropout_p = rng.choice([0.3, 0.4])
+    hidden = rng.choice([64, 128])
+    num_classes = rng.choice([3, 4])
+
+    buggy_code = f'''import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, TensorDataset
+
+torch.manual_seed(42)
+
+class TabularNet(nn.Module):
+    def __init__(self, input_dim, num_classes):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, {hidden}),
+            nn.BatchNorm1d({hidden}),
+            nn.ReLU(),
+            nn.Dropout(p={dropout_p}),
+            nn.Linear({hidden}, {hidden}),
+            nn.ReLU(),
+            nn.Dropout(p={dropout_p}),
+            nn.Linear({hidden}, num_classes),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+torch.manual_seed(42)
+N, D, C = 1000, 20, {num_classes}
+X_raw = torch.randn(N, D)
+true_weights = torch.randn(D, C)
+y_all = (X_raw @ true_weights).argmax(dim=1)
+
+# BUG 1: normalization computed on full dataset before split
+mean = X_raw.mean(dim=0)
+std = X_raw.std(dim=0) + 1e-8
+X_normalized = (X_raw - mean) / std
+
+split = int(0.8 * N)
+X_train, X_test = X_normalized[:split], X_normalized[split:]
+y_train, y_test = y_all[:split], y_all[split:]
+
+train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
+
+model = TabularNet(D, C)
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+criterion = nn.CrossEntropyLoss()
+
+for epoch in range(10):
+    model.train()
+    for xb, yb in train_loader:
+        optimizer.zero_grad()
+        loss = criterion(model(xb), yb)
+        loss.backward()
+        optimizer.step()
+    print(f"Epoch {{epoch+1}} complete")
+
+# BUG 2: model.eval() and torch.no_grad() missing
+test_preds = model(X_test).argmax(dim=1)
+accuracy = (test_preds == y_test).float().mean().item()
+print(f"Test accuracy: {{accuracy:.4f}}")
+print("Evaluation complete")
+print("Training finished")
+'''
+
+    error_output = (
+        "Script runs to completion with no errors.\n"
+        "Reported test accuracy: 0.9700 (varies slightly between runs)\n\n"
+        "This script has TWO silent bugs:\n\n"
+        "Bug 1 — Data leakage:\n"
+        "  Normalization statistics computed from entire dataset before train/test split.\n"
+        "  Test set distribution has contaminated preprocessing. Accuracy is inflated.\n\n"
+        "Bug 2 — Missing eval mode:\n"
+        f"  model.eval() not called before evaluation. Dropout(p={dropout_p}) "
+        "remains active causing slightly different predictions each run.\n\n"
+        "Fix BOTH. Fixed version should have lower but trustworthy accuracy "
+        "and produce identical results on repeated evaluation runs."
+    )
+
+    return BugScenario(
+        task_id=TASK_COMPOUND_LEAKAGE_EVAL,
+        task_description=(
+            "This PyTorch script runs cleanly and reports impressive metrics — but contains "
+            "TWO silent bugs that make the evaluation invalid. "
+            "There is a data leakage bug in preprocessing AND a missing eval mode bug. "
+            "Fix both so the evaluation is correct and deterministic."
+        ),
+        buggy_code=buggy_code,
+        error_output=error_output,
+        correct_bug_type="compound_leakage_eval",
+        solution_hint=f"fix 1: compute mean/std only from X_train after split; fix 2: add model.eval() and torch.no_grad() before evaluation",
+        num_bugs=2,
     )

server/grader.py

@@ -1,3 +1,4 @@
+import ast
 import subprocess
 import sys
 import tempfile
@@ -17,6 +18,7 @@ class GradeResult:
 
 
 def grade(action_bug_type: str, action_diagnosis: str, fixed_code: str, scenario: BugScenario) -> GradeResult:
+    # "other" skips bug type check — goes straight to execution-based scoring
     type_correct = _check_bug_type(action_bug_type, scenario.correct_bug_type)
     if not type_correct:
         return GradeResult(
@@ -29,7 +31,7 @@ def grade(action_bug_type: str, action_diagnosis: str, fixed_code: str, scenario
             execution_output="(code not executed — bug type was wrong)",
         )
 
-    exec_output, ran_ok = _run_code(fixed_code, timeout=30)
+    exec_output, ran_ok = _run_code(fixed_code, timeout=40)
 
     if not ran_ok:
         return GradeResult(
@@ -64,7 +66,7 @@ def grade(action_bug_type: str, action_diagnosis: str, fixed_code: str, scenario
             execution_output=exec_output,
         )
 
-    success, success_feedback = _check_success_signal(scenario.task_id, exec_output)
+    success, success_feedback = _check_success_signal(scenario.task_id, fixed_code, exec_output)
     if not success:
         return GradeResult(
             score=0.8,
@@ -79,8 +81,8 @@ def grade(action_bug_type: str, action_diagnosis: str, fixed_code: str, scenario
     return GradeResult(
         score=0.99,
         feedback=(
-            f"Perfect fix. Bug type correct, code runs cleanly, "
-            f"training completes, and success signal confirmed.\n"
+            "Perfect fix. Bug type correct, code runs cleanly, "
+            "training completes, and success signal confirmed.\n"
             f"Execution output:\n{exec_output}"
         ),
         execution_output=exec_output,
@@ -88,15 +90,44 @@ def grade(action_bug_type: str, action_diagnosis: str, fixed_code: str, scenario
 
 
 def _check_bug_type(submitted: str, correct: str) -> bool:
+    # "other" always passes — execution-based scoring handles it
     submitted_clean = submitted.strip().lower().replace(" ", "_").replace("-", "_")
+    if submitted_clean == "other":
+        return True
+
     correct_clean = correct.strip().lower()
     aliases = {
-        "shape_mismatch": {"shape_mismatch", "shape", "dimension", "size_mismatch", "linear", "matmul", "incompatible", "input_shape", "classifier"},
-        "training_collapse": {"training_collapse", "collapse", "nan", "diverge", "learning_rate", "loss_fn", "loss_function", "wrong_loss"},
-        "data_leakage": {"data_leakage", "leakage", "leak", "train_test_leak", "normalization", "preprocessing"},
-        "wrong_device": {"wrong_device", "device", "device_mismatch", "cuda", "cpu", "device_error"},
-        "gradient_not_zeroed": {"gradient_not_zeroed", "gradient", "zero_grad", "missing_zero_grad", "accumulate", "gradient_accumulation"},
-        "missing_eval_mode": {"missing_eval_mode", "eval_mode", "eval", "dropout", "batchnorm", "no_grad", "inference_mode"},
+        "shape_mismatch": {
+            "shape_mismatch", "shape", "dimension", "size_mismatch", "linear",
+            "matmul", "incompatible", "input_shape", "classifier",
+        },
+        "training_collapse": {
+            "training_collapse", "collapse", "nan", "diverge", "learning_rate",
+            "loss_fn", "loss_function", "wrong_loss",
+        },
+        "data_leakage": {
+            "data_leakage", "leakage", "leak", "train_test_leak",
+            "normalization", "preprocessing",
+        },
+        "wrong_device": {
+            "wrong_device", "device", "device_mismatch", "cuda", "cpu", "device_error",
+        },
+        "gradient_not_zeroed": {
+            "gradient_not_zeroed", "gradient", "zero_grad", "missing_zero_grad",
+            "accumulate", "gradient_accumulation",
+        },
+        "missing_eval_mode": {
+            "missing_eval_mode", "eval_mode", "eval", "dropout", "batchnorm",
+            "no_grad", "inference_mode",
+        },
+        "compound_shape_device": {
+            "compound_shape_device", "compound", "shape_device", "multiple",
+            "two_bugs", "shape_mismatch", "wrong_device", "shape", "device",
+        },
+        "compound_leakage_eval": {
+            "compound_leakage_eval", "compound", "leakage_eval", "multiple",
+            "two_bugs", "data_leakage", "missing_eval_mode", "leakage", "eval",
+        },
     }
     valid = aliases.get(correct_clean, {correct_clean})
     if submitted_clean in valid:
@@ -107,7 +138,7 @@ def _check_bug_type(submitted: str, correct: str) -> bool:
     return False
 
 
-def _run_code(code: str, timeout: int = 30) -> tuple[str, bool]:
+def _run_code(code: str, timeout: int = 40) -> tuple[str, bool]:
     with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8") as f:
         f.write(code)
         tmp_path = f.name
@@ -143,6 +174,13 @@ def _run_code(code: str, timeout: int = 30) -> tuple[str, bool]:
             pass
 
 
+def _run_code_twice(code: str) -> tuple[str, str, bool]:
+    """Run code twice and return both outputs. Used for determinism checks."""
+    out1, ok1 = _run_code(code, timeout=40)
+    out2, ok2 = _run_code(code, timeout=40)
+    return out1, out2, (ok1 and ok2)
+
+
 def _check_training_completed(output: str, task_id: str) -> bool:
     lower = output.lower()
     if "nan" in lower and "loss" in lower:
@@ -160,23 +198,65 @@ def _check_training_completed(output: str, task_id: str) -> bool:
     return any(m in lower for m in markers)
 
 
+def _zero_grad_before_backward_ast(code: str) -> bool:
+    """Use AST to verify optimizer.zero_grad() appears before loss.backward() in the loop."""
+    try:
+        tree = ast.parse(code)
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.For, ast.While)):
+                body = node.body
+                zero_grad_idx = None
+                backward_idx = None
+                for i, stmt in enumerate(body):
+                    stmt_str = ast.unparse(stmt) if hasattr(ast, 'unparse') else str(stmt)
+                    if "zero_grad" in stmt_str:
+                        zero_grad_idx = i
+                    if "backward" in stmt_str:
+                        backward_idx = i
+                if zero_grad_idx is not None and backward_idx is not None:
+                    if zero_grad_idx < backward_idx:
+                        return True
+        # Also check nested loops
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.For, ast.While)):
+                for child in ast.walk(node):
+                    if isinstance(child, (ast.For, ast.While)) and child is not node:
+                        body = child.body
+                        zero_grad_idx = None
+                        backward_idx = None
+                        for i, stmt in enumerate(body):
+                            stmt_str = ast.unparse(stmt) if hasattr(ast, 'unparse') else str(stmt)
+                            if "zero_grad" in stmt_str:
+                                zero_grad_idx = i
+                            if "backward" in stmt_str:
+                                backward_idx = i
+                        if zero_grad_idx is not None and backward_idx is not None:
+                            if zero_grad_idx < backward_idx:
+                                return True
+        return False
+    except Exception:
+        # AST parse failed — fall back to string check
+        return "optimizer.zero_grad()" in code
+
+
+def _extract_metric(output: str, pattern: str) -> Optional[float]:
+    match = re.search(pattern, output.lower())
+    if match:
+        try:
+            return float(match.group(1))
+        except ValueError:
+            return None
+    return None
+
+
 def _verify_fix(fixed_code: str, scenario: BugScenario, exec_output: str) -> tuple[bool, str]:
     task = scenario.task_id
 
     if task == "shape_mismatch":
-        hint = scenario.solution_hint
-        match = re.search(r"must be (\d+) not (\d+)", hint)
-        if match:
-            wrong_size = match.group(2)
-            # Check encoder final output matches classifier input
-            # Just verify code runs without shape error — already confirmed in stage 2/3
-            # Only fail if the EXACT original wrong Linear is unchanged
-            original_wrong = f"nn.Linear({wrong_size}, "
-            lines = fixed_code.split("\n")
-            classifier_lines = [l for l in lines if "classifier" in l.lower() and "nn.Linear" in l]
-            encoder_lines = [l for l in lines if "encoder" not in l.lower() and "nn.Linear" in l and "classifier" not in l.lower()]
-            # If code runs and trains (already verified), the shape is fixed
-            return True, ""
+        # Code runs and trains — shape is fixed by definition
+        # Just verify no shape error in output
+        if "cannot be multiplied" in exec_output.lower():
+            return False, "Shape mismatch error still present in execution output."
         return True, ""
 
     elif task == "training_collapse":
@@ -185,17 +265,24 @@ def _verify_fix(fixed_code: str, scenario: BugScenario, exec_output: str) -> tup
             return False, "Loss is still NaN in the fixed code output."
         hint = scenario.solution_hint
         if "learning rate" in hint or "lr" in hint.lower():
-            lr_matches = re.findall(r"lr\s*=\s*([\d.e\-+]+)", fixed_code)
+            lr_matches = re.findall(r"\blr\s*=\s*([\d.e\-+]+)", fixed_code)
             for lr_str in lr_matches:
                 try:
-                    lr_val = float(lr_str)
-                    if lr_val > 1.0:
-                        return False, f"Learning rate {lr_val} is still too large."
+                    if float(lr_str) > 1.0:
+                        return False, f"Learning rate {lr_str} is still too large."
                 except ValueError:
                     pass
         return True, ""
 
     elif task == "data_leakage":
+        # Behavioral check: accuracy should be in realistic range (not suspiciously high)
+        acc = _extract_metric(exec_output, r"accuracy[:\s]+([\d.]+)")
+        if acc is not None and acc > 0.97:
+            return False, f"Accuracy {acc:.4f} still suspiciously high — leakage may remain."
+        mse = _extract_metric(exec_output, r"mse[:\s]+([\d.]+)")
+        if mse is not None and mse < 0.04:
+            return False, f"MSE {mse:.4f} still suspiciously low — leakage may remain."
+        # Code check: no bad normalization pattern
         bad_patterns = [
             r"mean\s*=\s*[Xx]_raw\.mean",
             r"full_mean\s*=\s*[Xx]_raw\.mean",
@@ -203,55 +290,88 @@ def _verify_fix(fixed_code: str, scenario: BugScenario, exec_output: str) -> tup
         ]
         for pat in bad_patterns:
             if re.search(pat, fixed_code):
-                return False, "Normalization statistics still computed from full dataset before split."
-        good_patterns = [
-            r"train.*mean",
-            r"mean.*train",
-            r"X_train.*\.mean",
-            r"X_train_raw.*\.mean",
-        ]
-        has_good = any(re.search(p, fixed_code, re.IGNORECASE) for p in good_patterns)
-        if not has_good:
-            return False, "Could not confirm that normalization uses only training data statistics."
+                return False, "Normalization statistics still computed from full dataset."
         return True, ""
 
     elif task == "wrong_device":
-        # Fixed code must move data to device inside the loop
-        good_patterns = [
-            r"xb\s*=\s*xb\.to\(device\)",
-            r"xb\.to\(device\)",
-            r"\.to\(device\)",
-        ]
-        has_device_move = any(re.search(p, fixed_code) for p in good_patterns)
-        if not has_device_move:
-            return False, "Data tensors are not being moved to device inside the training loop."
-        # Must not crash with device error
         if "expected all tensors" in exec_output.lower():
             return False, "Device mismatch error still present in output."
+        good_patterns = [r"\.to\(device\)", r"\.to\('cpu'\)", r"\.to\(\"cpu\"\)"]
+        if not any(re.search(p, fixed_code) for p in good_patterns):
+            return False, "Data tensors not being moved to device inside the training loop."
         return True, ""
 
     elif task == "gradient_not_zeroed":
-        # Fixed code must have zero_grad before backward
-        if "optimizer.zero_grad()" not in fixed_code and "optim.zero_grad()" not in fixed_code:
-            return False, "optimizer.zero_grad() is still missing from the training loop."
-        lower_output = exec_output.lower()
-        if "nan" in lower_output and "loss" in lower_output:
+        if "nan" in exec_output.lower():
             return False, "Loss is still NaN — gradients may still be accumulating."
+        # AST-based check: zero_grad must appear before backward in loop body
+        if not _zero_grad_before_backward_ast(fixed_code):
+            return False, "optimizer.zero_grad() not found before loss.backward() in the loop."
         return True, ""
 
     elif task == "missing_eval_mode":
-        # Fixed code must have model.eval() before evaluation
+        # Behavioral check: run fixed code twice, outputs must be identical
+        out1, out2, both_ran = _run_code_twice(fixed_code)
+        if not both_ran:
+            return False, "Fixed code failed to run twice cleanly."
+        # Extract metrics from both runs
+        acc1 = _extract_metric(out1, r"accuracy[:\s]+([\d.]+)")
+        acc2 = _extract_metric(out2, r"accuracy[:\s]+([\d.]+)")
+        mse1 = _extract_metric(out1, r"mse[:\s]+([\d.]+)")
+        mse2 = _extract_metric(out2, r"mse[:\s]+([\d.]+)")
+        if acc1 is not None and acc2 is not None:
+            if abs(acc1 - acc2) > 0.02:
+                return False, f"Evaluation still non-deterministic: accuracy {acc1:.4f} vs {acc2:.4f} across two runs. model.eval() may be missing or ineffective."
+        if mse1 is not None and mse2 is not None:
+            if abs(mse1 - mse2) > 0.05:
+                return False, f"Evaluation still non-deterministic: MSE {mse1:.4f} vs {mse2:.4f} across two runs."
+        # Also check string presence as backup
         if "model.eval()" not in fixed_code:
             return False, "model.eval() is missing before the evaluation block."
-        # Should also have no_grad
-        if "torch.no_grad()" not in fixed_code and "no_grad" not in fixed_code:
-            return False, "torch.no_grad() context manager is missing during evaluation."
+        return True, ""
+
+    elif task == "compound_shape_device":
+        # Must fix BOTH: shape mismatch AND device placement
+        if "cannot be multiplied" in exec_output.lower():
+            return False, "Shape mismatch error still present — Bug 1 not fully fixed."
+        if "expected all tensors" in exec_output.lower():
+            return False, "Device mismatch error still present — Bug 2 not fully fixed."
+        good_device = any(re.search(p, fixed_code) for p in [r"\.to\(device\)", r"xb\.to\(", r"yb\.to\("])
+        if not good_device:
+            return False, "Data tensors not moved to device — Bug 2 not fixed."
+        return True, ""
+
+    elif task == "compound_leakage_eval":
+        # Must fix BOTH: data leakage AND missing eval mode
+        # Check 1: no bad normalization pattern
+        bad_patterns = [
+            r"mean\s*=\s*[Xx]_raw\.mean",
+            r"full_mean\s*=\s*[Xx]_raw\.mean",
+            r"[Xx]_raw\.mean\(dim=0\)",
+        ]
+        for pat in bad_patterns:
+            if re.search(pat, fixed_code):
+                return False, "Data leakage still present — normalization uses full dataset stats."
+        # Check 2: behavioral determinism test
+        out1, out2, both_ran = _run_code_twice(fixed_code)
+        if not both_ran:
+            return False, "Fixed code failed to run twice cleanly."
+        acc1 = _extract_metric(out1, r"accuracy[:\s]+([\d.]+)")
+        acc2 = _extract_metric(out2, r"accuracy[:\s]+([\d.]+)")
+        if acc1 is not None and acc2 is not None:
+            if abs(acc1 - acc2) > 0.02:
+                return False, f"Eval still non-deterministic: {acc1:.4f} vs {acc2:.4f}. model.eval() may be missing."
+        # Check 3: accuracy not suspiciously high (leakage check)
+        if acc1 is not None and acc1 > 0.97:
+            return False, f"Accuracy {acc1:.4f} still suspiciously high — data leakage may remain."
+        if "model.eval()" not in fixed_code:
+            return False, "model.eval() missing — Bug 2 not fixed."
         return True, ""
 
     return True, ""
 
 
-def _check_success_signal(task_id: str, output: str) -> tuple[bool, str]:
+def _check_success_signal(task_id: str, fixed_code: str, output: str) -> tuple[bool, str]:
     lower = output.lower()
 
     if task_id == "shape_mismatch":
@@ -259,39 +379,33 @@ def _check_success_signal(task_id: str, output: str) -> tuple[bool, str]:
         has_finished = "training finished" in lower or "complete" in lower
         if has_epoch and has_finished:
             return True, ""
-        return False, "Expected epoch completion messages and 'Training finished' not found."
+        return False, "Expected epoch logs and 'Training finished' not found."
 
     elif task_id == "training_collapse":
         if "nan" in lower:
-            return False, "Output still contains NaN loss values."
+            return False, "Output still contains NaN."
         loss_values = re.findall(r"loss[:\s]+([\d.]+)", lower)
         if len(loss_values) >= 2:
             first, last = float(loss_values[0]), float(loss_values[-1])
             if last < first * 0.95:
                 return True, ""
-            return False, f"Loss did not decrease: started at {first:.4f}, ended at {last:.4f}."
-        has_finished = "training finished" in lower
-        return has_finished, "Could not confirm loss decreased across epochs."
+            return False, f"Loss did not decrease: {first:.4f} → {last:.4f}."
+        return "training finished" in lower, "Could not confirm loss decreased."
 
     elif task_id == "data_leakage":
-        acc_match = re.search(r"accuracy[:\s]+([\d.]+)", lower)
-        if acc_match:
-            acc = float(acc_match.group(1))
-            if acc > 0.98:
-                return False, f"Reported accuracy {acc:.4f} is suspiciously high — leakage may still be present."
-        mse_match = re.search(r"mse[:\s]+([\d.]+)", lower)
-        if mse_match:
-            mse = float(mse_match.group(1))
-            if mse < 0.05:
-                return False, f"Reported MSE {mse:.4f} is suspiciously low — leakage may still be present."
-        has_finished = "training finished" in lower
-        return has_finished, "Training did not complete."
+        acc = _extract_metric(output, r"accuracy[:\s]+([\d.]+)")
+        if acc is not None and acc > 0.97:
+            return False, f"Accuracy {acc:.4f} suspiciously high — leakage may still be present."
+        mse = _extract_metric(output, r"mse[:\s]+([\d.]+)")
+        if mse is not None and mse < 0.04:
+            return False, f"MSE {mse:.4f} suspiciously low — leakage may still be present."
+        return "training finished" in lower, "Training did not complete."
 
     elif task_id == "wrong_device":
         has_epoch = any(f"epoch {i}" in lower for i in range(1, 4))
         has_finished = "training finished" in lower
-        no_device_error = "expected all tensors" not in lower
-        if has_epoch and has_finished and no_device_error:
+        no_error = "expected all tensors" not in lower
+        if has_epoch and has_finished and no_error:
             return True, ""
         return False, "Training did not complete cleanly or device error still present."
 
@@ -303,15 +417,31 @@ def _check_success_signal(task_id: str, output: str) -> tuple[bool, str]:
             first, last = float(loss_values[0]), float(loss_values[-1])
             if last < first * 0.9:
                 return True, ""
-            return False, f"Loss did not decrease sufficiently: {first:.4f} -> {last:.4f}."
-        has_finished = "training finished" in lower
-        return has_finished, "Could not confirm stable training."
+            return False, f"Loss did not decrease sufficiently: {first:.4f} → {last:.4f}."
+        return "training finished" in lower, "Could not confirm stable training."
 
     elif task_id == "missing_eval_mode":
-        has_accuracy = "accuracy" in lower or "complete" in lower
+        has_metric = "accuracy" in lower or "mse" in lower
+        has_finished = "training finished" in lower or "evaluation complete" in lower
+        if has_metric and has_finished:
+            return True, ""
+        return False, "Evaluation did not complete or metric not reported."
+
+    elif task_id == "compound_shape_device":
+        has_epoch = any(f"epoch {i}" in lower for i in range(1, 4))
+        has_finished = "training finished" in lower
+        no_errors = "cannot be multiplied" not in lower and "expected all tensors" not in lower
+        if has_epoch and has_finished and no_errors:
+            return True, ""
+        return False, "Training did not complete or one of the bugs is still present."
+
+    elif task_id == "compound_leakage_eval":
+        acc = _extract_metric(output, r"accuracy[:\s]+([\d.]+)")
         has_finished = "training finished" in lower or "evaluation complete" in lower
-        if has_accuracy and has_finished:
+        if acc is not None and acc > 0.97:
+            return False, f"Accuracy {acc:.4f} too high — data leakage may still be present."
+        if has_finished:
             return True, ""
-        return False, "Evaluation did not complete or accuracy not reported."
+        return False, "Training or evaluation did not complete."
 
     return True, ""

server/ml_debug_env_environment.py

@@ -11,15 +11,16 @@ from openenv.core.env_server.types import State
 from models import DebugAction, DebugObservation, DebugState
 from bug_generator import (
     get_scenario,
-    get_random_task,
     BugScenario,
+    ALL_TASKS,
     TASK_SHAPE_MISMATCH,
     TASK_TRAINING_COLLAPSE,
     TASK_DATA_LEAKAGE,
     TASK_WRONG_DEVICE,
     TASK_GRADIENT_NOT_ZEROED,
     TASK_MISSING_EVAL_MODE,
-    ALL_TASKS,
+    TASK_COMPOUND_SHAPE_DEVICE,
+    TASK_COMPOUND_LEAKAGE_EVAL,
 )
 from grader import grade, GradeResult
 
@@ -28,18 +29,23 @@ MAX_STEPS = 3
 
 class MlDebugEnvEnvironment(Environment):
     """
-    ML Debugging Environment — 6 tasks, easy → hard.
-
-    Tasks:
-      shape_mismatch      (easy)        — explicit crash, wrong linear layer size
-      training_collapse   (medium)      — NaN loss or wrong loss function
-      data_leakage        (hard)        — silent, evaluation is invalid
-      wrong_device        (medium)      — CPU/CUDA tensor mismatch, explicit crash
-      gradient_not_zeroed (medium-hard) — missing zero_grad, loss explodes
-      missing_eval_mode   (hard)        — no model.eval(), unreliable metrics
-
-    Episodes are single-step by default. MAX_STEPS=3 allows retry with
-    grader feedback fed back to the agent.
+    ML Debugging Environment — 8 tasks, easy → expert.
+
+    Single-bug tasks (6):
+      shape_mismatch      (easy)
+      training_collapse   (medium)
+      wrong_device        (medium)
+      gradient_not_zeroed (medium-hard)
+      data_leakage        (hard)
+      missing_eval_mode   (hard)
+
+    Compound tasks — TWO bugs per script (2):
+      compound_shape_device  (medium-hard) — shape mismatch + device mismatch
+      compound_leakage_eval  (expert)      — data leakage + missing eval mode
+
+    Graders are execution-based: fixed code is actually run in a subprocess.
+    Multi-turn episodes: agent gets up to 3 attempts with grader feedback.
+    bug_type="other" skips type check — goes straight to execution scoring.
     """
 
     SUPPORTS_CONCURRENT_SESSIONS = True
@@ -87,6 +93,7 @@ class MlDebugEnvEnvironment(Environment):
             grader_score=None,
             grader_feedback=None,
             step_number=0,
+            num_bugs=scenario.num_bugs,
             done=False,
             reward=None,
         )
@@ -113,10 +120,7 @@ class MlDebugEnvEnvironment(Environment):
         if result.score > self._state.current_score:
             self._state.current_score = result.score
 
-        done = (
-            result.score >= 0.95
-            or self._state.step_count >= MAX_STEPS
-        )
+        done = result.score >= 0.95 or self._state.step_count >= MAX_STEPS
 
         return DebugObservation(
             task_id=self._state.task_id,
@@ -127,6 +131,7 @@ class MlDebugEnvEnvironment(Environment):
             grader_score=result.score,
             grader_feedback=result.feedback,
             step_number=self._state.step_count,
+            num_bugs=self._current_scenario.num_bugs,
             done=done,
             reward=result.score,
         )
@@ -146,13 +151,10 @@ class MlDebugEnvEnvironment(Environment):
             name="ML Debugging Environment",
             description=(
                 "An RL environment where agents debug broken PyTorch training scripts. "
-                "Six tasks of increasing difficulty: shape mismatch (easy), "
-                "training collapse (medium), wrong device (medium), "
-                "gradient not zeroed (medium-hard), data leakage (hard), "
-                "and missing eval mode (hard). "
-                "Agents receive a buggy script and must return a corrected version. "
-                "The grader executes the fix and scores 0.01–0.99 with partial credit."
+                "Eight tasks: six single-bug (easy→hard) and two compound double-bug tasks (expert). "
+                "Graders execute fixed code in a subprocess — no shortcuts. "
+                "Multi-turn episodes with grader feedback. Accepts 'other' as bug_type for open-ended debugging."
             ),
-            version="2.0.0",
+            version="3.0.0",
             author="ml-debug-env",
         )