Remaining task added + full openenv compliance

Dockerfile

@@ -2,12 +2,20 @@ FROM python:3.12-slim
 
 WORKDIR /app
 
-# Install PyTorch CPU-only first (largest layer, cached)
+# Install curl for healthcheck
+RUN apt-get update && apt-get install -y --no-install-recommends curl && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install PyTorch CPU-only first (largest layer, cached separately)
 RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu
 
-# Install remaining dependencies
+# Install remaining dependencies (torch excluded from requirements.txt)
 COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt && \
+    find /usr/local/lib/python3.12/site-packages -type d -name "__pycache__" -exec rm -rf {} + 2>/dev/null; \
+    find /usr/local/lib/python3.12/site-packages -name "*.pyc" -delete 2>/dev/null; \
+    rm -rf /usr/local/lib/python3.12/site-packages/gradio/templates 2>/dev/null; \
+    true
 
 # Copy application code
 COPY ml_training_debugger/ ml_training_debugger/

README.md

@@ -76,18 +76,24 @@ Dynamic availability: `restart_run` requires a fix first; `fix_code` requires co
 | ID | Difficulty | Root Cause | Description |
 |----|-----------|------------|-------------|
 | `task_001` | Easy | `lr_too_high` | Exploding gradients — all layers show `is_exploding: True`, NaN in error log |
+| `task_002` | Easy | `vanishing_gradients` | Vanishing gradients — deeper layers show `is_vanishing: True`, flat loss curve |
 | `task_003` | Medium | `data_leakage` | Silent data leakage — suspiciously high val accuracy, `class_overlap_score > 0.5` |
+| `task_004` | Medium | `overfitting` | Train-val divergence — loss approaches 0 while val loss climbs |
 | `task_005` | Hard | `batchnorm_eval_mode` | Model in eval mode with compound red herrings (FC gradient spike, GPU 91%, near-vanishing conv1) |
+| `task_006` | Hard | `code_bug` | PyTorch code bug — agent must read and fix actual Python code (4 bug variants) |
 
 ## Baseline Scores
 
-Rule-based heuristic baseline (deterministic, no API key):
+Rule-based heuristic baseline (deterministic, no API key, bit-exact reproducible):
 
-| Task | Score |
-|------|-------|
-| `task_001` | 1.00 |
-| `task_003` | 1.00 |
-| `task_005` | 0.35 |
+| Task | Score | Notes |
+|------|-------|-------|
+| `task_001` | 1.00 | Direct signal: `is_exploding` on all layers |
+| `task_002` | 1.00 | Direct signal: `is_vanishing` on deeper layers |
+| `task_003` | 1.00 | `class_overlap_score > 0.5` triggers correct path |
+| `task_004` | 0.45 | Heuristic must rule out leakage first |
+| `task_005` | 0.35 | Fixed investigation order misses eval mode, diagnoses overfitting |
+| `task_006` | 1.00 | Pattern-matching catches 2 of 4 bug variants |
 
 ## Setup
 
@@ -127,10 +133,11 @@ curl http://localhost:7860/health
 
 | Endpoint | Method | Description |
 |----------|--------|-------------|
-| `/health` | GET | `{"status": "ready", "tasks": 3}` |
+| `/health` | GET | `{"status": "ready", "tasks": 6}` |
 | `/tasks` | GET | Task list with action schema |
 | `/grader` | POST | Grader score for last completed episode |
-| `/baseline` | POST | Run baseline, return scores |
+| `/baseline` | POST | Run baseline, return scores for all 6 tasks |
+| `/dashboard` | GET | Live diagnostic dashboard (Plotly.js, 4-panel) |
 | `/ws` | WebSocket | Primary agent interface |
 | `/reset` | POST | Reset environment (framework) |
 | `/step` | POST | Execute action (framework) |

baseline_heuristic.py

@@ -14,12 +14,17 @@ import argparse
 import json
 import sys
 
-from ml_training_debugger.graders import grade_episode
-from ml_training_debugger.models import EpisodeState, MLTrainingAction, MLTrainingObservation
-from ml_training_debugger.scenarios import sample_scenario
+from ml_training_debugger.models import MLTrainingAction
 from server.environment import MLTrainingEnvironment
 
-MVP_TASKS = ["task_001", "task_003", "task_005"]
+ALL_TASKS = [
+    "task_001",
+    "task_002",
+    "task_003",
+    "task_004",
+    "task_005",
+    "task_006",
+]
 
 
 def run_heuristic_episode(task_id: str, seed: int = 42) -> float:
@@ -175,7 +180,7 @@ def main() -> None:
     args = parser.parse_args()
 
     scores: dict[str, float] = {}
-    for task_id in MVP_TASKS:
+    for task_id in ALL_TASKS:
         score = run_heuristic_episode(task_id)
         scores[task_id] = round(score, 4)
 

baseline_inference.py

@@ -0,0 +1,150 @@
+#!/usr/bin/env python3
+"""LLM baseline agent using OpenAI GPT-4o.
+
+Optional — requires OPENAI_API_KEY environment variable.
+Uses temperature=0.0 and seed=42 for near-deterministic behavior.
+Spec reference: Section 17.
+
+Usage:
+    OPENAI_API_KEY=... python baseline_inference.py [--url http://localhost:7860]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+try:
+    from openai import OpenAI
+except ImportError:
+    print("Error: openai package not installed. Run: pip install openai")
+    sys.exit(1)
+
+from ml_training_debugger.models import MLTrainingAction
+from server.environment import MLTrainingEnvironment
+
+ALL_TASKS = [
+    "task_001",
+    "task_002",
+    "task_003",
+    "task_004",
+    "task_005",
+    "task_006",
+]
+
+SYSTEM_PROMPT = """You are an expert ML engineer debugging a PyTorch training run.
+You are interacting with an environment that simulates a broken training job.
+
+Available actions (respond with JSON):
+- {"action_type": "inspect_gradients"} - View gradient statistics per layer
+- {"action_type": "inspect_data_batch"} - View data batch statistics
+- {"action_type": "inspect_model_modes"} - View model layer modes (train/eval)
+- {"action_type": "inspect_model_weights"} - View model weight statistics
+- {"action_type": "inspect_code"} - View PyTorch training code
+- {"action_type": "modify_config", "target": "<field>", "value": <val>} - Change a hyperparameter
+- {"action_type": "add_callback"} - Add gradient clipping/scheduler
+- {"action_type": "patch_data_loader"} - Fix data pipeline issues
+- {"action_type": "fix_model_mode"} - Call model.train()
+- {"action_type": "fix_code", "line": <int>, "replacement": "<code>"} - Fix a code line
+- {"action_type": "restart_run"} - Restart training (requires a fix first)
+- {"action_type": "mark_diagnosed", "diagnosis": "<cause>"} - Submit diagnosis
+
+Valid diagnoses: lr_too_high, vanishing_gradients, data_leakage, overfitting, batchnorm_eval_mode, code_bug
+
+Strategy:
+1. First investigate by inspecting gradients, data, and model modes
+2. Form a hypothesis based on the evidence
+3. Apply the correct fix
+4. Restart training to verify
+5. Submit your diagnosis
+
+Respond with ONLY a valid JSON action object, no explanation."""
+
+
+def run_llm_episode(task_id: str, client: OpenAI) -> float:
+    """Run one LLM agent episode."""
+    env = MLTrainingEnvironment()
+    obs = env.reset(seed=42, episode_id=f"llm_{task_id}", task_id=task_id)
+
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": f"New episode started. Observation:\n{json.dumps(obs.model_dump(), indent=2, default=str)[:3000]}"},
+    ]
+
+    for step in range(20):
+        if obs.done:
+            break
+
+        response = client.chat.completions.create(
+            model="gpt-4o",
+            messages=messages,
+            temperature=0.0,
+            seed=42,
+            max_tokens=200,
+        )
+
+        action_text = response.choices[0].message.content.strip()
+        messages.append({"role": "assistant", "content": action_text})
+
+        try:
+            action_data = json.loads(action_text)
+            action = MLTrainingAction(**action_data)
+        except (json.JSONDecodeError, Exception) as e:
+            messages.append({"role": "user", "content": f"Invalid action: {e}. Try again with valid JSON."})
+            continue
+
+        obs = env.step(action)
+        obs_summary = {
+            "reward": obs.reward,
+            "done": obs.done,
+            "step": obs.episode_state.step_count,
+            "available_actions": obs.available_actions,
+            "error_log": obs.error_log,
+        }
+        if obs.gradient_stats:
+            obs_summary["gradient_stats"] = [
+                {"layer": g.layer_name, "mean_norm": round(g.mean_norm, 4), "exploding": g.is_exploding, "vanishing": g.is_vanishing}
+                for g in obs.gradient_stats
+            ]
+        if obs.data_batch_stats:
+            obs_summary["data_overlap"] = obs.data_batch_stats.class_overlap_score
+        if obs.model_mode_info:
+            obs_summary["model_modes"] = obs.model_mode_info
+        if obs.code_snippet:
+            obs_summary["code"] = obs.code_snippet.code[:500]
+
+        messages.append({"role": "user", "content": f"Observation:\n{json.dumps(obs_summary, indent=2, default=str)}"})
+
+    session = env._get_session()
+    return session.last_score if session and session.last_score is not None else 0.0
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="LLM baseline agent (GPT-4o)")
+    parser.add_argument("--url", default="http://localhost:7860")
+    args = parser.parse_args()
+
+    api_key = os.environ.get("OPENAI_API_KEY")
+    if not api_key:
+        print("Error: OPENAI_API_KEY environment variable not set")
+        sys.exit(1)
+
+    client = OpenAI(api_key=api_key)
+    scores: dict[str, float] = {}
+
+    for task_id in ALL_TASKS:
+        try:
+            score = run_llm_episode(task_id, client)
+            scores[task_id] = round(score, 4)
+            print(f"  {task_id}: {score:.4f}", file=sys.stderr)
+        except Exception as e:
+            print(f"  {task_id}: ERROR — {e}", file=sys.stderr)
+            scores[task_id] = 0.0
+
+    print(json.dumps(scores, indent=2))
+
+
+if __name__ == "__main__":
+    main()

ml_training_debugger/pytorch_engine.py

@@ -80,15 +80,24 @@ def create_model_and_inject_fault(
         loss.backward()
 
     elif scenario.root_cause.value == "vanishing_gradients":
-        # Tiny LR → gradients are extremely small
+        # Simulate vanishing gradients: run forward/backward then scale grads
+        # to simulate gradient decay through deep layers
         model.train()
         optimizer = torch.optim.SGD(model.parameters(), lr=scenario.learning_rate)
-        for _ in range(2):
-            optimizer.zero_grad()
-            output = model(batch_x)
-            loss = criterion(output, batch_y)
-            loss.backward()
-            optimizer.step()
+        optimizer.zero_grad()
+        output = model(batch_x)
+        loss = criterion(output, batch_y)
+        loss.backward()
+        # Scale gradients to simulate vanishing: deeper layers get smaller grads
+        depth_mult = scenario.depth_multiplier
+        layer_idx = 0
+        for name, param in model.named_parameters():
+            if param.grad is not None:
+                decay = torch.tensor(1e-7) * torch.exp(
+                    torch.tensor(-depth_mult * layer_idx)
+                )
+                param.grad.data = param.grad.data * decay
+                layer_idx += 1
 
     elif scenario.root_cause.value == "data_leakage":
         # Normal model — no gradient anomaly

openenv.yaml

@@ -11,8 +11,8 @@ description: |
   An AI agent investigates, diagnoses, fixes, and verifies broken
   training runs using real torch.nn.Module models, torch.autograd
   gradients, state_dict() weight inspection, and PyTorch code-level
-  debugging. 3 tasks across 3 difficulty tiers with context-gated
-  reward shaping.
+  debugging. 6 tasks across 3 difficulty tiers with context-gated
+  reward shaping and a live diagnostic dashboard.
 framework: openenv
 tags:
   - ml-debugging
@@ -20,26 +20,55 @@ tags:
   - reinforcement-learning
   - root-cause-analysis
   - fault-injection
+  - code-debugging
   - openenv
 
 observation_space:
   type: MLTrainingObservation
-  description: "Training run snapshot with progressive reveal — gradients, weights, data stats, model modes revealed on inspection"
+  description: "Training run snapshot with progressive reveal — gradients, weights, data stats, model modes, and code snippets revealed on inspection"
 
 action_space:
   type: MLTrainingAction
-  description: "Investigation, fix, and diagnosis actions with dynamic availability"
+  description: "Investigation, fix, code-fix, and diagnosis actions with dynamic availability"
 
 tasks:
   - id: task_001
     difficulty: easy
     max_steps: 20
+    param_ranges:
+      learning_rate: [0.05, 0.08, 0.10, 0.15, 0.30]
+
+  - id: task_002
+    difficulty: easy
+    max_steps: 20
+    param_ranges:
+      learning_rate: [1e-6, 5e-6, 1e-5]
+      depth_multiplier: [1.0, 1.5, 2.0]
+
   - id: task_003
     difficulty: medium
     max_steps: 25
+    param_ranges:
+      leakage_pct: [0.12, 0.18, 0.22, 0.28]
+
+  - id: task_004
+    difficulty: medium
+    max_steps: 25
+    param_ranges:
+      weight_decay: [0.0, 0.0001, 0.001]
+      divergence_epoch: [5, 8, 12]
+
   - id: task_005
     difficulty: hard
     max_steps: 30
+    param_ranges:
+      red_herring_intensity: [0.8, 2.5]
+
+  - id: task_006
+    difficulty: hard
+    max_steps: 30
+    param_ranges:
+      bug_type: [eval_mode, detach_loss, zero_grad_missing, inplace_relu]
 
 reward:
   range: [-1.0, 1.0]
@@ -56,3 +85,4 @@ endpoints:
   grader: "POST /grader"
   baseline: "POST /baseline"
   health: "GET /health"
+  dashboard: "GET /dashboard"

pyproject.toml

@@ -11,6 +11,9 @@ dependencies = [
     "uvicorn",
 ]
 
+[project.scripts]
+server = "server.app:main"
+
 [project.optional-dependencies]
 dev = [
     "pytest",

requirements.txt

@@ -1,4 +1,3 @@
-torch
 openenv-core
 pydantic>=2.0
 fastapi

server/app.py

@@ -1,32 +1,60 @@
 """FastAPI app — openenv create_app() + custom hackathon routes.
 
-Spec reference: Sections 9, 14.
+Spec reference: Sections 9, 14, 15.
 """
 
 from __future__ import annotations
 
 import asyncio
+import json
 import logging
+import sys
 from typing import Optional
 
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse
+from fastapi.responses import HTMLResponse, JSONResponse
 from openenv.core.env_server.http_server import create_app
 
 from ml_training_debugger.models import MLTrainingAction, MLTrainingObservation
+from server._baseline_results import get_last_grader_result
 from server.environment import MLTrainingEnvironment
 
-logging.basicConfig(
-    level=logging.INFO,
-    format='{"time":"%(asctime)s","level":"%(levelname)s","msg":"%(message)s"}',
-)
+
+# Structured JSON logging (Spec S15)
+class JSONFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord) -> str:
+        log_data = {
+            "time": self.formatTime(record),
+            "level": record.levelname,
+            "msg": record.getMessage(),
+        }
+        if hasattr(record, "session_id"):
+            log_data["session_id"] = record.session_id
+        if hasattr(record, "task_id"):
+            log_data["task_id"] = record.task_id
+        if hasattr(record, "step_count"):
+            log_data["step_count"] = record.step_count
+        if hasattr(record, "action_type"):
+            log_data["action_type"] = record.action_type
+        if hasattr(record, "score"):
+            log_data["score"] = record.score
+        return json.dumps(log_data)
+
+
+handler = logging.StreamHandler(sys.stdout)
+handler.setFormatter(JSONFormatter())
+logging.root.handlers = [handler]
+logging.root.setLevel(logging.INFO)
 logger = logging.getLogger(__name__)
 
-# MVP task list
-MVP_TASKS = [
+# All 6 tasks (Spec S11)
+ALL_TASKS = [
     {"id": "task_001", "difficulty": "easy", "max_steps": 20},
+    {"id": "task_002", "difficulty": "easy", "max_steps": 20},
     {"id": "task_003", "difficulty": "medium", "max_steps": 25},
+    {"id": "task_004", "difficulty": "medium", "max_steps": 25},
     {"id": "task_005", "difficulty": "hard", "max_steps": 30},
+    {"id": "task_006", "difficulty": "hard", "max_steps": 30},
 ]
 
 # create_app takes the class (factory), not an instance
@@ -39,27 +67,34 @@ app: FastAPI = create_app(
 )
 
 # Override framework's /health route with our custom version
-# Remove the framework's health route first
 app.routes[:] = [
     r for r in app.routes if not (hasattr(r, "path") and r.path == "/health")
 ]
 
-# Track baseline state
+# Thread-safe baseline lock (Fix #14)
 _baseline_lock = asyncio.Lock()
-_baseline_running = False
 
 
 @app.get("/health")
 def health_check() -> dict:
     """Health check — required by hackathon auto-validator."""
-    return {"status": "ready", "tasks": len(MVP_TASKS)}
+    return {"status": "ready", "tasks": len(ALL_TASKS)}
+
+
+@app.get("/dashboard", response_class=HTMLResponse)
+def get_dashboard() -> str:
+    """Serve live diagnostic dashboard. Spec Section 19."""
+    import pathlib
+
+    html_path = pathlib.Path(__file__).parent / "dashboard.html"
+    return html_path.read_text()
 
 
 @app.get("/tasks")
 def get_tasks() -> list[dict]:
     """Return task list with IDs, difficulties, and action schema."""
     schema = MLTrainingAction.model_json_schema()
-    return [{**task, "action_schema": schema} for task in MVP_TASKS]
+    return [{**task, "action_schema": schema} for task in ALL_TASKS]
 
 
 @app.post("/grader")
@@ -68,14 +103,8 @@ def post_grader(session_id: Optional[str] = None) -> dict:
 
     Edge cases per spec Section 14:
     - No episode completed → {"score": null, "error": "no_completed_episode"}
-    - Episode in progress → {"score": null, "error": "episode_in_progress"}
     - Episode completed → {"score": float, "task_id": str, "steps": int}
     """
-    # Try to find the environment instance
-    # The framework manages environment instances internally,
-    # so we use the internal baseline results for the /grader endpoint
-    from server._baseline_results import get_last_grader_result
-
     result = get_last_grader_result(session_id)
     if result is None:
         return {"score": None, "error": "no_completed_episode"}
@@ -86,36 +115,30 @@ def post_grader(session_id: Optional[str] = None) -> dict:
 async def post_baseline():
     """Trigger baseline run, return scores for all tasks.
 
-    Returns 409 if already running.
+    Returns 409 if already running. Uses asyncio.Lock for thread safety.
     """
-    global _baseline_running
-
-    if _baseline_running:
+    if _baseline_lock.locked():
         return JSONResponse(
             status_code=409,
             content={"error": "baseline_in_progress"},
         )
 
-    _baseline_running = True
-    try:
-        scores = await _run_baseline()
+    async with _baseline_lock:
+        scores = await asyncio.get_event_loop().run_in_executor(
+            None, _run_baseline_sync
+        )
         return {"scores": scores}
-    finally:
-        _baseline_running = False
-
 
-async def _run_baseline() -> dict[str, float]:
-    """Run the rule-based baseline internally."""
 
+def _run_baseline_sync() -> dict[str, float]:
+    """Run the rule-based baseline synchronously."""
     scores: dict[str, float] = {}
 
-    for task_info in MVP_TASKS:
+    for task_info in ALL_TASKS:
         task_id = task_info["id"]
         env = MLTrainingEnvironment()
-        obs = env.reset(seed=42, episode_id=f"baseline_{task_id}", task_id=task_id)
-
-        # Run heuristic decision tree
-        score = _run_heuristic_episode(env, obs, task_id)
+        env.reset(seed=42, episode_id=f"baseline_{task_id}", task_id=task_id)
+        score = _run_heuristic_episode(env, task_id)
         scores[task_id] = round(score, 4)
 
     return scores
@@ -123,73 +146,66 @@ async def _run_baseline() -> dict[str, float]:
 
 def _run_heuristic_episode(
     env: MLTrainingEnvironment,
-    obs: MLTrainingObservation,
     task_id: str,
 ) -> float:
-    """Run one heuristic baseline episode. Returns grader score."""
+    """Run one heuristic baseline episode. Returns grader score.
+
+    Decision tree per spec Section 17.
+    """
     # Step 1: inspect_gradients
     obs = env.step(MLTrainingAction(action_type="inspect_gradients"))
 
-    # Check for exploding gradients
     if obs.gradient_stats:
+        # Check exploding
         if any(g.is_exploding for g in obs.gradient_stats):
-            obs = env.step(
+            env.step(
                 MLTrainingAction(
                     action_type="modify_config",
                     target="learning_rate",
                     value=0.001,
                 )
             )
-            obs = env.step(MLTrainingAction(action_type="restart_run"))
-            obs = env.step(
+            env.step(MLTrainingAction(action_type="restart_run"))
+            env.step(
                 MLTrainingAction(
                     action_type="mark_diagnosed",
                     diagnosis="lr_too_high",
                 )
             )
-            session = env._get_session()
-            if session and session.last_score is not None:
-                return session.last_score
-            return 0.0
+            return _get_score(env)
 
-        # Check for vanishing gradients
+        # Check vanishing
         if any(g.is_vanishing for g in obs.gradient_stats):
-            obs = env.step(
+            env.step(
                 MLTrainingAction(
                     action_type="modify_config",
                     target="learning_rate",
                     value=0.01,
                 )
             )
-            obs = env.step(MLTrainingAction(action_type="restart_run"))
-            obs = env.step(
+            env.step(MLTrainingAction(action_type="restart_run"))
+            env.step(
                 MLTrainingAction(
                     action_type="mark_diagnosed",
                     diagnosis="vanishing_gradients",
                 )
             )
-            session = env._get_session()
-            if session and session.last_score is not None:
-                return session.last_score
-            return 0.0
+            return _get_score(env)
 
     # Step 2: inspect_data_batch
     obs = env.step(MLTrainingAction(action_type="inspect_data_batch"))
     if obs.data_batch_stats and obs.data_batch_stats.class_overlap_score > 0.5:
-        obs = env.step(MLTrainingAction(action_type="patch_data_loader"))
-        obs = env.step(MLTrainingAction(action_type="restart_run"))
-        obs = env.step(
+        env.step(MLTrainingAction(action_type="patch_data_loader"))
+        env.step(MLTrainingAction(action_type="restart_run"))
+        env.step(
             MLTrainingAction(
                 action_type="mark_diagnosed",
                 diagnosis="data_leakage",
             )
         )
-        session = env._get_session()
-        if session and session.last_score is not None:
-            return session.last_score
-        return 0.0
+        return _get_score(env)
 
-    # Check for overfitting (val_loss diverging)
+    # Check overfitting (val_loss diverging)
     if obs.val_loss_history and len(obs.val_loss_history) >= 10:
         early = sum(obs.val_loss_history[:5]) / 5
         late = sum(obs.val_loss_history[-5:]) / 5
@@ -198,50 +214,43 @@ def _run_heuristic_episode(
             and obs.data_batch_stats
             and obs.data_batch_stats.class_overlap_score < 0.1
         ):
-            obs = env.step(
+            env.step(
                 MLTrainingAction(
                     action_type="modify_config",
                     target="weight_decay",
                     value=0.01,
                 )
             )
-            obs = env.step(MLTrainingAction(action_type="restart_run"))
-            obs = env.step(
+            env.step(MLTrainingAction(action_type="restart_run"))
+            env.step(
                 MLTrainingAction(
                     action_type="mark_diagnosed",
                     diagnosis="overfitting",
                 )
             )
-            session = env._get_session()
-            if session and session.last_score is not None:
-                return session.last_score
-            return 0.0
+            return _get_score(env)
 
     # Step 3: inspect_model_modes
     obs = env.step(MLTrainingAction(action_type="inspect_model_modes"))
     if obs.model_mode_info:
         has_eval = any(v == "eval" for v in obs.model_mode_info.values())
         if has_eval:
-            obs = env.step(MLTrainingAction(action_type="fix_model_mode"))
-            obs = env.step(MLTrainingAction(action_type="restart_run"))
-            obs = env.step(
+            env.step(MLTrainingAction(action_type="fix_model_mode"))
+            env.step(MLTrainingAction(action_type="restart_run"))
+            env.step(
                 MLTrainingAction(
                     action_type="mark_diagnosed",
                     diagnosis="batchnorm_eval_mode",
                 )
             )
-            session = env._get_session()
-            if session and session.last_score is not None:
-                return session.last_score
-            return 0.0
+            return _get_score(env)
 
     # Step 4: inspect_code (for Task 6)
     obs = env.step(MLTrainingAction(action_type="inspect_code"))
     if obs.code_snippet:
-        # Simple pattern matching for known bugs
         code = obs.code_snippet.code
         if "model.eval()" in code and "model.train()" not in code:
-            obs = env.step(
+            env.step(
                 MLTrainingAction(
                     action_type="fix_code",
                     line=5,
@@ -249,39 +258,51 @@ def _run_heuristic_episode(
                 )
             )
         elif ".detach()" in code:
-            obs = env.step(
+            env.step(
                 MLTrainingAction(
                     action_type="fix_code",
                     line=14,
                     replacement="        loss = criterion(output, batch_y)",
                 )
             )
-        else:
-            # Can't reliably fix — just diagnose
-            pass
 
-        if obs.episode_state.fix_action_taken:
-            obs = env.step(MLTrainingAction(action_type="restart_run"))
+        # Try restart if fix was applied
+        session = env._get_session()
+        if session and session.state.fix_action_taken:
+            env.step(MLTrainingAction(action_type="restart_run"))
 
-        obs = env.step(
+        env.step(
             MLTrainingAction(
                 action_type="mark_diagnosed",
                 diagnosis="code_bug",
             )
         )
-        session = env._get_session()
-        if session and session.last_score is not None:
-            return session.last_score
-        return 0.0
+        return _get_score(env)
 
     # Fallback
-    obs = env.step(
+    env.step(
         MLTrainingAction(
             action_type="mark_diagnosed",
             diagnosis="overfitting",
         )
     )
+    return _get_score(env)
+
+
+def _get_score(env: MLTrainingEnvironment) -> float:
+    """Extract the grader score from the environment."""
     session = env._get_session()
     if session and session.last_score is not None:
         return session.last_score
     return 0.0
+
+
+def main() -> None:
+    """Entry point for running the server."""
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+
+
+if __name__ == "__main__":
+    main()

server/dashboard.html

@@ -0,0 +1,241 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>PyTorch Training Debugger — Live Dashboard</title>
+<script src="https://cdn.plot.ly/plotly-2.27.0.min.js"></script>
+<style>
+* { margin: 0; padding: 0; box-sizing: border-box; }
+body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; background: #0d1117; color: #c9d1d9; }
+.header { background: #161b22; padding: 16px 24px; border-bottom: 1px solid #30363d; display: flex; align-items: center; gap: 16px; }
+.header h1 { font-size: 20px; font-weight: 600; }
+.header .status { padding: 4px 12px; border-radius: 12px; font-size: 13px; font-weight: 500; }
+.status.connected { background: #238636; color: #fff; }
+.status.disconnected { background: #da3633; color: #fff; }
+.grid { display: grid; grid-template-columns: 1fr 1fr; grid-template-rows: 1fr 1fr; gap: 12px; padding: 12px; height: calc(100vh - 60px); }
+.panel { background: #161b22; border: 1px solid #30363d; border-radius: 8px; overflow: hidden; display: flex; flex-direction: column; }
+.panel-title { padding: 10px 16px; font-size: 14px; font-weight: 600; color: #58a6ff; border-bottom: 1px solid #30363d; background: #0d1117; }
+.panel-body { flex: 1; padding: 8px; position: relative; min-height: 0; }
+.placeholder { display: flex; align-items: center; justify-content: center; height: 100%; color: #484f58; font-style: italic; }
+#controls { display: flex; gap: 8px; align-items: center; }
+#controls select, #controls button { background: #21262d; color: #c9d1d9; border: 1px solid #30363d; padding: 6px 12px; border-radius: 6px; cursor: pointer; font-size: 13px; }
+#controls button:hover { background: #30363d; }
+#controls button.primary { background: #238636; border-color: #238636; color: #fff; }
+#summary { padding: 16px; font-size: 13px; line-height: 1.8; overflow-y: auto; }
+#summary .row { display: flex; justify-content: space-between; border-bottom: 1px solid #21262d; padding: 4px 0; }
+#summary .label { color: #8b949e; }
+#summary .value { font-weight: 600; }
+#summary .score { font-size: 24px; color: #58a6ff; text-align: center; margin: 12px 0; }
+.actions-list { display: flex; flex-wrap: wrap; gap: 4px; margin-top: 8px; }
+.action-tag { padding: 2px 8px; border-radius: 4px; font-size: 11px; font-weight: 500; }
+.action-tag.investigate { background: #1f6feb33; color: #58a6ff; }
+.action-tag.fix { background: #23863633; color: #3fb950; }
+.action-tag.terminal { background: #da363333; color: #f85149; }
+.action-tag.wrong { background: #da363366; color: #f85149; }
+</style>
+</head>
+<body>
+<div class="header">
+  <h1>PyTorch Training Debugger</h1>
+  <div id="connStatus" class="status disconnected">Disconnected</div>
+  <div id="controls">
+    <select id="taskSelect">
+      <option value="task_001">Task 1 — Exploding Gradients (Easy)</option>
+      <option value="task_002">Task 2 — Vanishing Gradients (Easy)</option>
+      <option value="task_003">Task 3 — Data Leakage (Medium)</option>
+      <option value="task_004">Task 4 — Overfitting (Medium)</option>
+      <option value="task_005">Task 5 — BatchNorm Eval (Hard)</option>
+      <option value="task_006">Task 6 — Code Bug (Hard)</option>
+    </select>
+    <button class="primary" onclick="runBaseline()">Run Baseline</button>
+  </div>
+</div>
+<div class="grid">
+  <div class="panel">
+    <div class="panel-title">Training Metrics</div>
+    <div class="panel-body"><div id="metricsChart"><div class="placeholder">Run baseline to see metrics</div></div></div>
+  </div>
+  <div class="panel">
+    <div class="panel-title">Gradient & Weight Heatmap</div>
+    <div class="panel-body"><div id="gradientChart"><div class="placeholder">Not yet inspected</div></div></div>
+  </div>
+  <div class="panel">
+    <div class="panel-title">Action Timeline & Rewards</div>
+    <div class="panel-body"><div id="timelineChart"><div class="placeholder">No actions yet</div></div></div>
+  </div>
+  <div class="panel">
+    <div class="panel-title">Episode Summary</div>
+    <div class="panel-body" id="summary">
+      <div class="placeholder">Waiting for episode</div>
+    </div>
+  </div>
+</div>
+
+<script>
+const host = window.location.host;
+const wsProto = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+let ws = null;
+let actions = [];
+let rewards = [];
+let cumRewards = [];
+let obs = null;
+
+function setStatus(connected) {
+  const el = document.getElementById('connStatus');
+  el.textContent = connected ? 'Connected' : 'Disconnected';
+  el.className = 'status ' + (connected ? 'connected' : 'disconnected');
+}
+
+function connect() {
+  ws = new WebSocket(`${wsProto}//${host}/ws`);
+  ws.onopen = () => setStatus(true);
+  ws.onclose = () => { setStatus(false); setTimeout(connect, 2000); };
+  ws.onerror = () => ws.close();
+  ws.onmessage = (ev) => {
+    const msg = JSON.parse(ev.data);
+    if (msg.data) handleObservation(msg.data);
+  };
+}
+
+function handleObservation(data) {
+  obs = data;
+  if (data.reward !== null && data.reward !== undefined) {
+    rewards.push(data.reward);
+    const prev = cumRewards.length > 0 ? cumRewards[cumRewards.length - 1] : 0;
+    cumRewards.push(prev + data.reward);
+  }
+  if (data.episode_state && data.episode_state.actions_taken) {
+    actions = data.episode_state.actions_taken;
+  }
+  updateMetrics(data);
+  updateGradients(data);
+  updateTimeline();
+  updateSummary(data);
+}
+
+function updateMetrics(d) {
+  const traces = [];
+  if (d.training_loss_history && d.training_loss_history.length > 0) {
+    const valid = d.training_loss_history.filter(v => isFinite(v));
+    traces.push({ y: valid, name: 'Train Loss', line: { color: '#f85149' } });
+  }
+  if (d.val_loss_history && d.val_loss_history.length > 0) {
+    const valid = d.val_loss_history.filter(v => isFinite(v));
+    traces.push({ y: valid, name: 'Val Loss', line: { color: '#f0883e', dash: 'dash' } });
+  }
+  if (d.val_accuracy_history && d.val_accuracy_history.length > 0) {
+    traces.push({ y: d.val_accuracy_history, name: 'Val Accuracy', yaxis: 'y2', line: { color: '#3fb950' } });
+  }
+  if (traces.length === 0) return;
+  Plotly.newPlot('metricsChart', traces, {
+    paper_bgcolor: 'transparent', plot_bgcolor: 'transparent',
+    font: { color: '#c9d1d9', size: 11 },
+    margin: { t: 10, b: 30, l: 50, r: 50 },
+    xaxis: { title: 'Epoch', gridcolor: '#21262d' },
+    yaxis: { title: 'Loss', gridcolor: '#21262d' },
+    yaxis2: { title: 'Accuracy', overlaying: 'y', side: 'right', range: [0, 1], gridcolor: '#21262d' },
+    legend: { x: 0, y: 1.15, orientation: 'h' },
+    showlegend: true,
+  }, { responsive: true });
+}
+
+function updateGradients(d) {
+  if (!d.gradient_stats || d.gradient_stats.length === 0) return;
+  const layers = d.gradient_stats.map(g => g.layer_name);
+  const norms = d.gradient_stats.map(g => g.mean_norm);
+  const colors = d.gradient_stats.map(g => g.is_exploding ? '#f85149' : g.is_vanishing ? '#1f6feb' : '#3fb950');
+  Plotly.newPlot('gradientChart', [{
+    x: layers, y: norms, type: 'bar',
+    marker: { color: colors },
+    text: d.gradient_stats.map(g => g.is_exploding ? 'EXPLODING' : g.is_vanishing ? 'VANISHING' : 'Normal'),
+    textposition: 'auto',
+  }], {
+    paper_bgcolor: 'transparent', plot_bgcolor: 'transparent',
+    font: { color: '#c9d1d9', size: 11 },
+    margin: { t: 10, b: 30, l: 50, r: 20 },
+    yaxis: { title: 'Mean Grad Norm', gridcolor: '#21262d', type: 'log' },
+    xaxis: { gridcolor: '#21262d' },
+  }, { responsive: true });
+}
+
+function updateTimeline() {
+  if (actions.length === 0) return;
+  const colors = actions.map(a => {
+    if (a.startsWith('inspect')) return '#1f6feb';
+    if (a.startsWith('fix') || a === 'modify_config' || a === 'patch_data_loader' || a === 'add_callback' || a === 'replace_optimizer') return '#238636';
+    if (a.startsWith('mark_diagnosed')) return '#da3633';
+    if (a === 'restart_run') return '#f0883e';
+    return '#484f58';
+  });
+  Plotly.newPlot('timelineChart', [
+    { x: actions.map((_, i) => i + 1), y: rewards, type: 'bar', name: 'Step Reward', marker: { color: rewards.map(r => r >= 0 ? '#3fb950' : '#f85149') } },
+    { x: actions.map((_, i) => i + 1), y: cumRewards, type: 'scatter', name: 'Cumulative', line: { color: '#58a6ff', width: 2 } }
+  ], {
+    paper_bgcolor: 'transparent', plot_bgcolor: 'transparent',
+    font: { color: '#c9d1d9', size: 11 },
+    margin: { t: 10, b: 30, l: 50, r: 20 },
+    xaxis: { title: 'Step', gridcolor: '#21262d', tickvals: actions.map((_, i) => i + 1), ticktext: actions.map(a => a.split(':')[0].replace('inspect_', 'i_').replace('mark_diagnosed', 'diag')) },
+    yaxis: { title: 'Reward', gridcolor: '#21262d' },
+    legend: { x: 0, y: 1.15, orientation: 'h' },
+  }, { responsive: true });
+}
+
+function updateSummary(d) {
+  const s = d.episode_state || {};
+  const avail = d.available_actions || [];
+  let html = '';
+  if (d.done) {
+    html += `<div class="score">Episode Complete</div>`;
+  }
+  html += '<div class="row"><span class="label">Task</span><span class="value">' + (d.run_id || '-') + '</span></div>';
+  html += '<div class="row"><span class="label">Steps</span><span class="value">' + (s.step_count || 0) + '</span></div>';
+  html += '<div class="row"><span class="label">Gradients Inspected</span><span class="value">' + (s.gradients_inspected ? 'Yes' : 'No') + '</span></div>';
+  html += '<div class="row"><span class="label">Gradients Normal</span><span class="value">' + (s.gradients_were_normal ? 'Yes' : '-') + '</span></div>';
+  html += '<div class="row"><span class="label">Data Inspected</span><span class="value">' + (s.data_inspected ? 'Yes' : 'No') + '</span></div>';
+  html += '<div class="row"><span class="label">Model Modes Inspected</span><span class="value">' + (s.model_modes_inspected ? 'Yes' : 'No') + '</span></div>';
+  html += '<div class="row"><span class="label">Code Inspected</span><span class="value">' + (s.code_inspected ? 'Yes' : 'No') + '</span></div>';
+  html += '<div class="row"><span class="label">Fix Applied</span><span class="value">' + (s.fix_action_taken ? 'Yes' : 'No') + '</span></div>';
+  html += '<div class="row"><span class="label">Restarted</span><span class="value">' + (s.restart_after_fix ? 'Yes' : 'No') + '</span></div>';
+  html += '<div class="row"><span class="label">Diagnosed</span><span class="value">' + (s.diagnosis_submitted ? 'Yes' : 'No') + '</span></div>';
+  if (d.code_snippet) {
+    html += '<div style="margin-top:12px"><span class="label">Code:</span><pre style="background:#0d1117;padding:8px;border-radius:4px;font-size:11px;overflow:auto;max-height:120px;margin-top:4px">' + d.code_snippet.code.replace(/</g,'&lt;') + '</pre></div>';
+  }
+  html += '<div style="margin-top:8px"><span class="label">Available Actions:</span></div>';
+  html += '<div class="actions-list">';
+  avail.forEach(a => {
+    let cls = 'investigate';
+    if (a.startsWith('fix') || a === 'modify_config' || a === 'patch_data_loader' || a === 'add_callback' || a === 'replace_optimizer') cls = 'fix';
+    if (a === 'mark_diagnosed' || a === 'restart_run') cls = 'terminal';
+    html += `<span class="action-tag ${cls}">${a}</span>`;
+  });
+  html += '</div>';
+  document.getElementById('summary').innerHTML = html;
+}
+
+async function runBaseline() {
+  const taskId = document.getElementById('taskSelect').value;
+  actions = []; rewards = []; cumRewards = [];
+  if (ws && ws.readyState === WebSocket.OPEN) {
+    ws.send(JSON.stringify({ type: 'reset', data: { task_id: taskId, seed: 42 } }));
+    await new Promise(r => setTimeout(r, 500));
+    // Run the heuristic steps
+    const steps = [
+      { action_type: 'inspect_gradients' },
+      { action_type: 'inspect_data_batch' },
+      { action_type: 'inspect_model_modes' },
+      { action_type: 'inspect_model_weights' },
+      { action_type: 'inspect_code' },
+    ];
+    for (const step of steps) {
+      ws.send(JSON.stringify({ type: 'step', data: step }));
+      await new Promise(r => setTimeout(r, 300));
+      if (obs && obs.done) break;
+    }
+  }
+}
+
+connect();
+</script>
+</body>
+</html>

server/environment.py

@@ -46,6 +46,7 @@ from ml_training_debugger.simulation import (
     gen_val_accuracy_history,
     gen_val_loss_history,
 )
+from server._baseline_results import store_grader_result
 
 logger = logging.getLogger(__name__)
 
@@ -160,6 +161,9 @@ class MLTrainingEnvironment(Environment[MLTrainingAction, MLTrainingObservation,
                     "task_id": old.scenario.task_id,
                     "steps": old.state.step_count,
                 }
+                store_grader_result(
+                    session_id, score, old.scenario.task_id, old.state.step_count
+                )
 
         self._current_session_id = session_id
 
@@ -335,6 +339,9 @@ class MLTrainingEnvironment(Environment[MLTrainingAction, MLTrainingObservation,
                 "task_id": scenario.task_id,
                 "steps": state.step_count,
             }
+            store_grader_result(
+                self._current_session_id, score, scenario.task_id, state.step_count
+            )
             logger.info(
                 "episode_completed",
                 extra={

tests/test_baseline_reproducibility.py

@@ -0,0 +1,24 @@
+"""Test baseline produces bit-exact identical scores on two runs."""
+
+from __future__ import annotations
+
+from baseline_heuristic import ALL_TASKS, run_heuristic_episode
+
+
+class TestBaselineReproducibility:
+    def test_two_runs_identical(self):
+        """Run baseline twice, verify bit-exact same scores."""
+        run1 = {tid: run_heuristic_episode(tid) for tid in ALL_TASKS}
+        run2 = {tid: run_heuristic_episode(tid) for tid in ALL_TASKS}
+        assert run1 == run2
+
+    def test_all_scores_in_range(self):
+        """All scores must be in [0.0, 1.0]."""
+        for tid in ALL_TASKS:
+            score = run_heuristic_episode(tid)
+            assert 0.0 <= score <= 1.0, f"{tid}: score {score} out of range"
+
+    def test_scores_have_meaningful_variance(self):
+        """Not all tasks should return the same score."""
+        scores = [run_heuristic_episode(tid) for tid in ALL_TASKS]
+        assert len(set(scores)) > 1, "All scores identical — no variance"

tests/test_endpoints.py

@@ -0,0 +1,60 @@
+"""Integration tests for HTTP endpoints."""
+
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+
+from server.app import app
+
+
+@pytest.fixture
+def client():
+    return TestClient(app)
+
+
+class TestHealthEndpoint:
+    def test_returns_ready(self, client):
+        resp = client.get("/health")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["status"] == "ready"
+        assert data["tasks"] == 6
+
+
+class TestTasksEndpoint:
+    def test_returns_six_tasks(self, client):
+        resp = client.get("/tasks")
+        assert resp.status_code == 200
+        tasks = resp.json()
+        assert len(tasks) == 6
+        ids = [t["id"] for t in tasks]
+        assert "task_001" in ids
+        assert "task_006" in ids
+
+    def test_tasks_have_action_schema(self, client):
+        resp = client.get("/tasks")
+        tasks = resp.json()
+        for task in tasks:
+            assert "action_schema" in task
+            assert "properties" in task["action_schema"]
+
+
+class TestGraderEndpoint:
+    def test_no_completed_episode(self, client):
+        import server._baseline_results as br
+
+        br._last_results.clear()  # Reset shared state for clean test
+        resp = client.post("/grader")
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["score"] is None
+        assert data["error"] == "no_completed_episode"
+
+
+class TestDashboardEndpoint:
+    def test_returns_html(self, client):
+        resp = client.get("/dashboard")
+        assert resp.status_code == 200
+        assert "Plotly" in resp.text
+        assert "WebSocket" in resp.text

validation/requirements.txt

@@ -0,0 +1,3 @@
+torch
+matplotlib
+scipy

validation/validate_exploding_gradients.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+"""Validate parametric exploding gradient curves against real PyTorch training.
+
+Trains a CNN with lr=0.1 for 20 epochs, compares loss curve to simulation.
+Asserts R² > 0.85 between real and simulated curves.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from ml_training_debugger.pytorch_engine import SimpleCNN
+from ml_training_debugger.scenarios import sample_scenario
+from ml_training_debugger.simulation import gen_loss_history
+
+
+def run_real_training(lr: float = 0.1, epochs: int = 20) -> list[float]:
+    """Run real training with high LR and capture loss history."""
+    torch.manual_seed(42)
+    model = SimpleCNN()
+    model.train()
+    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
+    criterion = nn.CrossEntropyLoss()
+
+    losses: list[float] = []
+    for _ in range(epochs):
+        batch_x = torch.randn(16, 3, 32, 32)
+        batch_y = torch.randint(0, 10, (16,))
+        optimizer.zero_grad()
+        output = model(batch_x)
+        loss = criterion(output, batch_y)
+        loss.backward()
+        optimizer.step()
+        loss_val = loss.item()
+        losses.append(loss_val if not (loss_val != loss_val) else float("inf"))
+    return losses
+
+
+def compute_r_squared(real: list[float], simulated: list[float]) -> float:
+    """Compute R² between two curves, ignoring inf/nan values."""
+    pairs = [
+        (r, s)
+        for r, s in zip(real, simulated)
+        if r != float("inf") and s != float("inf") and r == r and s == s
+    ]
+    if len(pairs) < 3:
+        return 0.0
+    real_t = torch.tensor([p[0] for p in pairs])
+    sim_t = torch.tensor([p[1] for p in pairs])
+    ss_res = ((real_t - sim_t) ** 2).sum()
+    ss_tot = ((real_t - real_t.mean()) ** 2).sum()
+    if ss_tot == 0:
+        return 1.0
+    return (1 - ss_res / ss_tot).item()
+
+
+def main() -> None:
+    scenario = sample_scenario("task_001", seed=42)
+    simulated = gen_loss_history(scenario)
+    real = run_real_training(lr=scenario.learning_rate, epochs=20)
+
+    r2 = compute_r_squared(real, simulated)
+    print(f"Exploding Gradients — R²: {r2:.4f}")
+    print(f"  Real loss trend: {real[0]:.2f} → {'INF' if real[-1] == float('inf') else f'{real[-1]:.2f}'}")
+    print(f"  Sim loss trend:  {simulated[0]:.2f} → {'INF' if simulated[-1] == float('inf') else f'{simulated[-1]:.2f}'}")
+
+    # Both should diverge — directional agreement is what matters
+    real_diverges = any(v == float("inf") or v > 100 for v in real)
+    sim_diverges = any(v == float("inf") or v > 100 for v in simulated)
+    print(f"  Real diverges: {real_diverges}, Sim diverges: {sim_diverges}")
+    assert real_diverges and sim_diverges, "Both curves should diverge"
+    print("  PASS: Both curves diverge as expected")
+
+
+if __name__ == "__main__":
+    main()