Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,20 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN pip install --no-cache-dir \
+    torch==2.2.0+cpu torchvision==0.17.0+cpu \
+    --index-url https://download.pytorch.org/whl/cpu
+
+RUN apt-get update && apt-get install -y curl && rm -rf /var/lib/apt/lists/*
+
+COPY server/requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 7860
+ENV PYTHONUNBUFFERED=1
+ENV ENABLE_WEB_INTERFACE=false
+
+CMD ["python", "-m", "uvicorn", "server.app:app", "--host", "0.0.0.0", "--port", "7860"]

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -1,10 +1,52 @@
 ---
-title: Whipstudio
-emoji: 📊
-colorFrom: indigo
-colorTo: pink
+title: WhipStudio Env
+emoji: 🤖
+colorFrom: blue
+colorTo: green
 sdk: docker
-pinned: false
+app_port: 7860
+base_path: /ui
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ML Debug Environment
+
+An OpenEnv-compatible RL environment where agents debug broken PyTorch training scripts.
+
+## Environment Description
+The agent receives a broken Python training script and must return a corrected version.
+Five tasks simulate real ML production bugs with increasing complexity.
+
+## Action Space
+- fixed_code (str, required): Complete corrected Python script
+- explanation (str, optional): Description of bugs found
+- attempt_number (int, 1-3): Which fix attempt this is
+
+## Observation Space
+- task_id: Which task (task1/task2/task3)
+- task_description: Plain English instructions
+- buggy_code: The broken script
+- error_log: stdout+stderr from previous attempt
+- last_reward: Score from previous attempt (0.0 on first step)
+- metrics: {exit_code, elapsed_seconds, timed_out, step, best_reward_so_far}
+
+## Reward Function
+Continuous score 0.0–1.0. Partial credit for every improvement.
+See `server/tasks/graders.py` for per-task scoring logic.
+
+## Tasks
+| Task | Difficulty | Bug Type |
+|------|-----------|----------|
+| task1 | Easy | Wrong optimizer order + bad learning rate |
+| task2 | Medium | Silent NaN from log(0) numerical instability |
+| task3 | Hard | OOM memory leak + train/val data leakage |
+| task4 | Medium | Wrong loss function |
+| task5 | Medium | Frozen backbone |
+
+## Setup
+```bash
+pip install openenv-core
+uvicorn server.app:app --host 0.0.0.0 --port 7860
+```
+
+## Endpoints
+POST /reset, POST /step, GET /state, GET /tasks, POST /grader, GET /baseline

__init__.py

@@ -0,0 +1,6 @@
+"""ML Debug OpenEnv package."""
+
+from .models import MLDebugAction, MLDebugObservation
+from .client import MLDebugEnv
+
+__all__ = ["MLDebugAction", "MLDebugObservation", "MLDebugEnv"]

baseline_agent.py

@@ -0,0 +1,185 @@
+import asyncio
+import os
+
+from dotenv import load_dotenv
+load_dotenv(override=True)
+
+import httpx
+
+SYSTEM_PROMPT = """
+You are an expert PyTorch debugging agent.
+You receive a broken training script and must fix ALL bugs in it.
+Rules:
+- Return ONLY the complete corrected Python code, nothing else.
+- No markdown, no backticks, no explanation text.
+- The script must print losses in format: LOSSES:[v1, v2, ...]
+- For task3, also print: VAL_ACCS:[v1,...] and FINAL_LOSS:X.XX
+- Keep all torch.manual_seed() calls intact.
+""".strip()
+
+
+def get_model():
+    from smolagents import InferenceClientModel
+
+    hf_token = os.environ.get("HF_TOKEN")
+    if not hf_token:
+        raise RuntimeError(
+            "HF_TOKEN is not set. Set HF_TOKEN to run /baseline with InferenceClientModel."
+        )
+
+    return InferenceClientModel(
+        model_id="Qwen/Qwen2.5-Coder-32B-Instruct",
+        token=hf_token,
+    )
+
+
+def _generate_fixed_code(model, prompt: str) -> str:
+    def _extract_text(response) -> str:
+        if isinstance(response, str):
+            return response
+
+        if hasattr(response, "content"):
+            content = getattr(response, "content")
+            if isinstance(content, str):
+                return content
+            if isinstance(content, list):
+                chunks = []
+                for item in content:
+                    if isinstance(item, str):
+                        chunks.append(item)
+                    elif isinstance(item, dict):
+                        text = item.get("text") or item.get("content")
+                        if text:
+                            chunks.append(str(text))
+                if chunks:
+                    return "\n".join(chunks)
+
+        if isinstance(response, dict):
+            text = response.get("content") or response.get("text")
+            if isinstance(text, str):
+                return text
+
+        return str(response)
+
+    if hasattr(model, "generate"):
+        generate = getattr(model, "generate")
+        messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": prompt},
+        ]
+        try:
+            return _extract_text(generate(messages=messages))
+        except TypeError:
+            return _extract_text(generate(messages))
+
+    if callable(model):
+        try:
+            return _extract_text(model(prompt, system_prompt=SYSTEM_PROMPT))
+        except TypeError:
+            return _extract_text(model(prompt))
+
+    raise AttributeError("Model does not support callable() or generate() inference APIs")
+
+
+async def run_single_task(task_id: str, env_url: str = "http://localhost:7860") -> float:
+    """Backwards-compatible wrapper that returns just the score."""
+    result = await run_single_task_detailed(task_id, env_url)
+    return result["score"]
+
+
+async def run_single_task_detailed(task_id: str, env_url: str = "http://localhost:7860") -> dict:
+    """Run the baseline agent on a single task. Returns detailed results."""
+    model = get_model()
+    timeout = httpx.Timeout(900.0, connect=10.0)
+
+    attempts_log = []
+
+    async with httpx.AsyncClient(timeout=timeout) as client:
+        reset_resp = await client.post(f"{env_url}/reset", json={"task_id": task_id})
+        reset_resp.raise_for_status()
+        obs = reset_resp.json().get("observation", reset_resp.json())
+
+        best_reward = 0.0
+        best_code = ""
+        best_output = ""
+
+        for attempt in range(1, 4):
+            prompt = f"""
+Task: {obs.get('task_description', '')}
+Buggy code:
+{obs.get('buggy_code', '')}
+Previous execution output (if any):
+{obs.get('error_log', 'None')}
+Previous score: {obs.get('last_reward', 0.0)}
+""".strip()
+
+            fixed_code = _generate_fixed_code(model, prompt)
+            if "```python" in fixed_code:
+                fixed_code = fixed_code.split("```python", 1)[1].split("```", 1)[0].strip()
+            elif "```" in fixed_code:
+                fixed_code = fixed_code.split("```", 1)[1].split("```", 1)[0].strip()
+
+            step_payload = {
+                "action": {
+                    "fixed_code": fixed_code,
+                    "attempt_number": attempt,
+                }
+            }
+            step_resp = await client.post(f"{env_url}/step", json=step_payload)
+            if step_resp.status_code == 422:
+                step_resp = await client.post(
+                    f"{env_url}/step",
+                    json={
+                        "fixed_code": fixed_code,
+                        "attempt_number": attempt,
+                    },
+                )
+            step_resp.raise_for_status()
+            result = step_resp.json()
+
+            reward = float(result.get("reward", 0.0) or 0.0)
+            obs = result.get("observation", obs)
+            output_log = obs.get("error_log", "") if isinstance(obs, dict) else ""
+
+            attempts_log.append({
+                "attempt": attempt,
+                "code": fixed_code,
+                "output": output_log[:3000],
+                "reward": reward,
+            })
+
+            if reward > best_reward:
+                best_reward = reward
+                best_code = fixed_code
+                best_output = output_log
+
+            if result.get("done") or reward >= 0.95:
+                break
+
+    return {
+        "score": best_reward,
+        "fixed_code": best_code,
+        "output": best_output[:3000],
+        "attempts": attempts_log,
+    }
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--env-url", default="http://localhost:7860")
+    args = parser.parse_args()
+
+    async def main():
+        scores = {}
+        for tid in ["task1", "task2", "task3"]:
+            try:
+                s = await asyncio.wait_for(run_single_task(tid, args.env_url), timeout=95.0)
+            except TimeoutError:
+                s = 0.0
+            scores[tid] = round(s, 4)
+            print(f"{tid}: {s:.4f}")
+        print(f"Average: {sum(scores.values()) / 3:.4f}")
+
+    asyncio.run(main())

client.py

@@ -0,0 +1,42 @@
+from typing import Dict
+
+from openenv.core import EnvClient
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+
+from .models import MLDebugAction, MLDebugObservation
+
+
+class MLDebugEnv(EnvClient[MLDebugAction, MLDebugObservation, State]):
+    def _step_payload(self, action: MLDebugAction) -> Dict:
+        return {
+            "fixed_code": action.fixed_code,
+            "explanation": action.explanation,
+            "attempt_number": action.attempt_number,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[MLDebugObservation]:
+        obs_data = payload.get("observation", {})
+        observation = MLDebugObservation(
+            task_id=obs_data.get("task_id", "task1"),
+            task_description=obs_data.get("task_description", ""),
+            buggy_code=obs_data.get("buggy_code", ""),
+            error_log=obs_data.get("error_log", ""),
+            last_reward=obs_data.get("last_reward", 0.0),
+            metrics=obs_data.get("metrics", {}),
+            done=payload.get("done", False),
+            reward=payload.get("reward"),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> State:
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

gradio_app.py

@@ -0,0 +1,756 @@
+"""
+WhipStudio — ML Debug Arena
+A polished Gradio UI for the ML Debugging RL environment.
+Provides code editing, loss curve visualization, diff views, and episode history.
+"""
+import difflib
+import json
+import math
+import re
+from typing import Any
+
+import gradio as gr
+import httpx
+import os
+
+DEFAULT_BASE_URL = os.environ.get("BASE_URL", "http://localhost:7860")
+
+# ── Task metadata ──────────────────────────────────────────────────────────
+
+TASK_INFO = {
+    "task1": {
+        "name": "Broken Training Loop",
+        "difficulty": "🟢 Easy",
+        "description": "Fix optimizer order + learning rate bugs in a linear classifier.",
+        "hints": "Look at optimizer.step() / loss.backward() order and the learning rate.",
+    },
+    "task2": {
+        "name": "Silent NaN Loss",
+        "difficulty": "🟡 Medium",
+        "description": "Fix numerical instability causing NaN loss from log(0).",
+        "hints": "The loss computation uses torch.log() without clamping — pred can be 0.",
+    },
+    "task3": {
+        "name": "OOM + Data Leakage",
+        "difficulty": "🔴 Hard",
+        "description": "Fix memory leak (graph accumulation) AND train/val data leakage.",
+        "hints": "Two bugs: total_loss accumulates graph, and augmentation is applied before split.",
+    },
+    "task4": {
+        "name": "Wrong Loss Function",
+        "difficulty": "🟡 Medium",
+        "description": "Multi-label classification incorrectly using CrossEntropyLoss. Fix loss and eval.",
+        "hints": "Use BCEWithLogitsLoss for multi-label. Ensure predictions are multi-hot.",
+    },
+    "task5": {
+        "name": "Frozen Backbone",
+        "difficulty": "🟡 Medium",
+        "description": "Backbone frozen but its parameters are passed to the optimizer.",
+        "hints": "Unfreeze backend or only pass head parameters to Adam.",
+    },
+}
+
+# ── Theme ──────────────────────────────────────────────────────────────────
+
+CUSTOM_CSS = """
+/* Dark arena theme */
+.gradio-container { max-width: 1400px !important; }
+
+/* Score display */
+.score-high { color: #22c55e !important; font-weight: 700; }
+.score-med  { color: #eab308 !important; font-weight: 700; }
+.score-low  { color: #ef4444 !important; font-weight: 700; }
+
+/* Task badges */
+.task-badge { display: inline-block; padding: 2px 8px; border-radius: 4px; font-size: 0.85em; }
+.badge-easy { background: #16a34a22; color: #22c55e; border: 1px solid #22c55e44; }
+.badge-medium { background: #ca8a0422; color: #eab308; border: 1px solid #eab30844; }
+.badge-hard { background: #dc262622; color: #ef4444; border: 1px solid #ef444444; }
+
+/* Diff highlighting */
+.diff-add { color: #22c55e; background: #22c55e11; }
+.diff-del { color: #ef4444; background: #ef444411; }
+
+/* Score bar */
+.score-bar-container { position: relative; height: 28px; background: #1e293b; border-radius: 6px; overflow: hidden; margin: 8px 0; }
+.score-bar-fill { height: 100%; border-radius: 6px; transition: width 0.6s ease-in-out; }
+.score-bar-label { position: absolute; right: 8px; top: 4px; font-weight: 600; font-size: 0.9em; }
+
+/* Episode step indicators */
+.step-indicator { display: inline-flex; align-items: center; gap: 6px; padding: 4px 12px; border-radius: 16px; font-size: 0.85em; margin: 2px; }
+.step-done { background: #16a34a22; border: 1px solid #22c55e44; color: #22c55e; }
+.step-active { background: #3b82f622; border: 1px solid #3b82f644; color: #60a5fa; }
+.step-pending { background: #334155; border: 1px solid #47556966; color: #94a3b8; }
+
+/* Header styling */
+.arena-header { text-align: center; padding: 12px 0; }
+.arena-header h1 { margin: 0; font-size: 1.8em; }
+.arena-header p { margin: 4px 0 0 0; color: #94a3b8; }
+"""
+
+
+# ── API helpers ────────────────────────────────────────────────────────────
+
+def _api(base_url: str, method: str, path: str, payload: dict | None = None) -> dict:
+    """Call the WhipStudio API and return parsed JSON."""
+    base_url = (base_url or DEFAULT_BASE_URL).strip().rstrip("/")
+    url = f"{base_url}{path}"
+    try:
+        with httpx.Client(timeout=90.0) as client:
+            if method == "GET":
+                resp = client.get(url)
+            else:
+                resp = client.post(url, json=payload or {})
+        resp.raise_for_status()
+        return resp.json()
+    except Exception as exc:
+        return {"error": f"{exc.__class__.__name__}: {exc}"}
+
+
+def _parse_losses_from_log(log: str) -> list[float]:
+    """Extract LOSSES:[...] from stdout."""
+    match = re.search(r"LOSSES:\[([^\]]+)\]", log)
+    if not match:
+        return []
+    try:
+        return [float(x.strip()) for x in match.group(1).split(",")]
+    except Exception:
+        return []
+
+
+def _parse_val_accs_from_log(log: str) -> list[float]:
+    """Extract VAL_ACCS:[...] from stdout."""
+    match = re.search(r"VAL_ACCS:\[([^\]]+)\]", log)
+    if not match:
+        return []
+    try:
+        return [float(x.strip()) for x in match.group(1).split(",")]
+    except Exception:
+        return []
+
+
+def _score_color(score: float) -> str:
+    if score >= 0.7:
+        return "#22c55e"
+    if score >= 0.4:
+        return "#eab308"
+    return "#ef4444"
+
+
+def _score_html(score: float) -> str:
+    pct = int(score * 100)
+    color = _score_color(score)
+    return f"""
+<div style="text-align:center; margin: 8px 0;">
+    <div style="font-size: 2.4em; font-weight: 700; color: {color};">{score:.2f}</div>
+    <div style="position:relative; height:24px; background:#1e293b; border-radius:6px; overflow:hidden; margin:8px auto; max-width:280px;">
+        <div style="height:100%; width:{pct}%; background:linear-gradient(90deg, {color}88, {color}); border-radius:6px; transition:width 0.6s ease;"></div>
+    </div>
+    <div style="color:#94a3b8; font-size:0.85em;">{pct}% complete</div>
+</div>"""
+
+
+def _diff_html(original: str, fixed: str) -> str:
+    """Generate an HTML diff view between original and fixed code."""
+    orig_lines = original.strip().splitlines()
+    fixed_lines = fixed.strip().splitlines()
+    diff = difflib.unified_diff(orig_lines, fixed_lines, lineterm="", n=3)
+    lines = []
+    for line in diff:
+        if line.startswith("+++") or line.startswith("---"):
+            continue
+        if line.startswith("@@"):
+            lines.append(f'<div style="color:#60a5fa;margin:8px 0 2px 0;font-size:0.85em;">{line}</div>')
+        elif line.startswith("+"):
+            lines.append(f'<div style="color:#22c55e;background:#22c55e0d;padding:1px 6px;font-family:monospace;font-size:0.85em;">+ {line[1:]}</div>')
+        elif line.startswith("-"):
+            lines.append(f'<div style="color:#ef4444;background:#ef44440d;padding:1px 6px;font-family:monospace;font-size:0.85em;">- {line[1:]}</div>')
+        else:
+            lines.append(f'<div style="color:#94a3b8;padding:1px 6px;font-family:monospace;font-size:0.85em;">  {line}</div>')
+    if not lines:
+        return '<div style="color:#94a3b8;text-align:center;padding:20px;">No changes detected</div>'
+    return '<div style="background:#0f172a;border-radius:8px;padding:12px;overflow-x:auto;max-height:500px;overflow-y:auto;">' + "\n".join(lines) + "</div>"
+
+
+def _step_timeline_html(trajectory: list[dict], current_step: int, max_steps: int = 3) -> str:
+    """Render step timeline as HTML."""
+    items = []
+    for i in range(1, max_steps + 1):
+        entry = next((t for t in trajectory if t.get("step") == i), None)
+        if entry:
+            r = entry["reward"]
+            color = _score_color(r)
+            items.append(
+                f'<span class="step-indicator step-done" style="border-color:{color}44;background:{color}11;color:{color};">'
+                f'Step {i} → {r:.2f}</span>'
+            )
+        elif i == current_step + 1:
+            items.append(f'<span class="step-indicator step-active">Step {i} ▶</span>')
+        else:
+            items.append(f'<span class="step-indicator step-pending">Step {i}</span>')
+
+    return '<div style="display:flex;gap:8px;justify-content:center;flex-wrap:wrap;padding:8px 0;">' + "".join(items) + "</div>"
+
+
+def _loss_plot(losses: list[float], title: str = "Loss Curve"):
+    """Generate a matplotlib figure for loss curve."""
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(5, 3))
+    fig.patch.set_facecolor("#0f172a")
+    ax.set_facecolor("#1e293b")
+
+    if losses:
+        valid_losses = [(i, l) for i, l in enumerate(losses) if not (math.isnan(l) or math.isinf(l))]
+        nan_steps = [i for i, l in enumerate(losses) if math.isnan(l) or math.isinf(l)]
+
+        if valid_losses:
+            steps, vals = zip(*valid_losses)
+            ax.plot(steps, vals, color="#60a5fa", linewidth=2, marker="o", markersize=3, zorder=3)
+            ax.fill_between(steps, vals, alpha=0.15, color="#60a5fa")
+
+        if nan_steps:
+            ax.scatter(nan_steps, [max(v for _, v in valid_losses) if valid_losses else 1.0] * len(nan_steps),
+                      color="#ef4444", marker="x", s=60, zorder=4, label="NaN/Inf")
+            ax.legend(facecolor="#1e293b", edgecolor="#334155", labelcolor="#94a3b8")
+
+    ax.set_xlabel("Step", color="#94a3b8", fontsize=9)
+    ax.set_ylabel("Loss", color="#94a3b8", fontsize=9)
+    ax.set_title(title, color="#e2e8f0", fontsize=11, fontweight="bold")
+    ax.tick_params(colors="#64748b", labelsize=8)
+    for spine in ax.spines.values():
+        spine.set_color("#334155")
+    ax.grid(True, alpha=0.15, color="#475569")
+
+    fig.tight_layout()
+    return fig
+
+
+def _val_acc_plot(accs: list[float]):
+    """Generate a matplotlib figure for validation accuracy."""
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(5, 3))
+    fig.patch.set_facecolor("#0f172a")
+    ax.set_facecolor("#1e293b")
+
+    if accs:
+        epochs = list(range(1, len(accs) + 1))
+        ax.plot(epochs, accs, color="#a78bfa", linewidth=2, marker="o", markersize=3)
+        ax.fill_between(epochs, accs, alpha=0.15, color="#a78bfa")
+        ax.axhline(y=0.7, color="#22c55e", linestyle="--", alpha=0.5, label="Target (0.70)")
+        ax.legend(facecolor="#1e293b", edgecolor="#334155", labelcolor="#94a3b8")
+
+    ax.set_xlabel("Epoch", color="#94a3b8", fontsize=9)
+    ax.set_ylabel("Accuracy", color="#94a3b8", fontsize=9)
+    ax.set_title("Validation Accuracy", color="#e2e8f0", fontsize=11, fontweight="bold")
+    ax.set_ylim(0, 1.05)
+    ax.tick_params(colors="#64748b", labelsize=8)
+    for spine in ax.spines.values():
+        spine.set_color("#334155")
+    ax.grid(True, alpha=0.15, color="#475569")
+
+    fig.tight_layout()
+    return fig
+
+
+# ── Episode state ──────────────────────────────────────────────────────────
+
+class EpisodeState:
+    """Track episode state across Gradio interactions."""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.task_id = ""
+        self.buggy_code = ""
+        self.task_description = ""
+        self.step = 0
+        self.best_reward = 0.0
+        self.last_reward = 0.0
+        self.error_log = ""
+        self.trajectory: list[dict] = []
+        self.done = False
+        self.last_fixed_code = ""
+
+
+_state = EpisodeState()
+
+
+# ── Action handlers ────────────────────────────────────────────────────────
+
+def do_reset(base_url: str, task_id: str):
+    """Reset the environment for a given task."""
+    _state.reset()
+    data = _api(base_url, "POST", "/reset", {"task_id": task_id})
+    if "error" in data:
+        return (
+            f"❌ Error: {data['error']}",      # status
+            "",                                  # code editor
+            "",                                  # task desc
+            _score_html(0.0),                    # score
+            None,                                # loss plot
+            None,                                # acc plot
+            "",                                  # diff
+            _step_timeline_html([], 0),          # timeline
+            "",                                  # error log
+        )
+
+    obs = data.get("observation", data)
+    _state.task_id = obs.get("task_id", task_id)
+    _state.buggy_code = obs.get("buggy_code", "")
+    _state.task_description = obs.get("task_description", "")
+    _state.step = 0
+    _state.done = False
+
+    info = TASK_INFO.get(task_id, {})
+    task_md = f"""### {info.get('name', task_id)}  {info.get('difficulty', '')}
+
+{_state.task_description.strip()}
+
+**💡 Hint:** {info.get('hints', 'No hints available.')}
+"""
+
+    return (
+        f"✅ Episode started — {info.get('name', task_id)}",
+        _state.buggy_code.strip(),
+        task_md,
+        _score_html(0.0),
+        _loss_plot([], "Loss Curve — Submit a fix to see results"),
+        None,
+        '<div style="color:#94a3b8;text-align:center;padding:20px;">Submit a fix to see the diff</div>',
+        _step_timeline_html([], 0),
+        "",
+    )
+
+
+def do_step(base_url: str, fixed_code: str):
+    """Submit a fix attempt."""
+    if _state.done:
+        return (
+            "⚠️ Episode is done. Reset to start a new one.",
+            _score_html(_state.best_reward),
+            None, None, "", _step_timeline_html(_state.trajectory, _state.step), "",
+        )
+
+    if not fixed_code or not fixed_code.strip():
+        return (
+            "⚠️ Please enter code before submitting.",
+            _score_html(_state.last_reward),
+            None, None, "", _step_timeline_html(_state.trajectory, _state.step), "",
+        )
+
+    _state.step += 1
+    _state.last_fixed_code = fixed_code
+
+    payload = {"action": {"fixed_code": fixed_code, "attempt_number": _state.step}}
+    data = _api(base_url, "POST", "/step", payload)
+
+    if "error" in data:
+        _state.step -= 1
+        return (
+            f"❌ Error: {data['error']}",
+            _score_html(_state.last_reward),
+            None, None, "", _step_timeline_html(_state.trajectory, _state.step), "",
+        )
+
+    reward = float(data.get("reward", 0.0) or 0.0)
+    _state.last_reward = reward
+    _state.best_reward = max(_state.best_reward, reward)
+    _state.done = data.get("done", False)
+
+    obs = data.get("observation", {})
+    _state.error_log = obs.get("error_log", "")
+    metrics = obs.get("metrics", {})
+
+    _state.trajectory.append({
+        "step": _state.step,
+        "reward": reward,
+        "best_reward": _state.best_reward,
+        "metrics": metrics,
+        "done": _state.done,
+    })
+
+    # Parse outputs for visualization
+    losses = _parse_losses_from_log(_state.error_log)
+    val_accs = _parse_val_accs_from_log(_state.error_log)
+
+    loss_fig = _loss_plot(losses, f"Loss Curve — Step {_state.step}")
+    acc_fig = _val_acc_plot(val_accs) if val_accs else None
+
+    diff = _diff_html(_state.buggy_code, fixed_code)
+    timeline = _step_timeline_html(_state.trajectory, _state.step)
+
+    if reward >= 0.95:
+        emoji = "🎯"
+    elif reward >= 0.7:
+        emoji = "✅"
+    elif len(_state.trajectory) > 1 and reward > _state.trajectory[-2]["reward"]:
+        emoji = "📈"
+    else:
+        emoji = "⚠️"
+    done_msg = " — Episode complete!" if _state.done else ""
+    status = f"{emoji} Step {_state.step}/3 — Reward: {reward:.2f} (Best: {_state.best_reward:.2f}){done_msg}"
+
+    error_display = _state.error_log if _state.error_log else "No errors — code ran successfully."
+
+    return (status, _score_html(_state.best_reward), loss_fig, acc_fig, diff, timeline, error_display)
+
+
+def do_run_baseline(base_url: str, task_id: str):
+    """Run the baseline agent on a single task."""
+    # First reset
+    reset_result = do_reset(base_url, task_id)
+    yield reset_result + ("🤖 Resetting environment...",)
+
+    # Call baseline endpoint
+    data = _api(base_url, "GET", "/baseline")
+    if "error" in data:
+        yield reset_result + (f"❌ Baseline error: {data['error']}",)
+        return
+
+    scores = data.get("baseline_scores", {})
+    avg = data.get("average", 0.0)
+
+    results_md = "### 🤖 Baseline Agent Results\n\n"
+    results_md += "| Task | Score |\n|---|---|\n"
+    for tid in ["task1", "task2", "task3", "task4", "task5"]:
+        s = scores.get(tid, 0.0)
+        emoji = "🎯" if s >= 0.9 else ("✅" if s >= 0.7 else ("📈" if s >= 0.4 else "⚠️"))
+        results_md += f"| {tid} | {emoji} {s:.4f} |\n"
+    results_md += f"\n**Average: {avg:.4f}**"
+
+    yield reset_result + (results_md,)
+
+
+def load_current_state(base_url: str):
+    """Fetch and format current environment state for UI display."""
+    data = _api(base_url, "GET", "/state")
+    if "error" in data:
+        summary = "⚠️ Could not fetch current state. Start or reset an episode first, then try again."
+        return summary, {"error": data["error"]}
+
+    state_obj = data.get("state", data)
+    if not isinstance(state_obj, dict):
+        return "⚠️ State endpoint returned an unexpected response format.", {"raw": data}
+
+    done = bool(state_obj.get("done", False))
+    step = state_obj.get("step", 0)
+    task_id = state_obj.get("task_id", "-")
+    reward = state_obj.get("last_reward", state_obj.get("reward", 0.0))
+    summary = (
+        f"**Task:** {task_id}  |  **Step:** {step}  |  "
+        f"**Done:** {'yes' if done else 'no'}  |  **Last Reward:** {reward}"
+    )
+    return summary, state_obj
+
+
+# ── Build the UI ───────────────────────────────────────────────────────────
+
+def build_ui() -> gr.Blocks:
+    theme = gr.themes.Soft(
+        primary_hue=gr.themes.colors.blue,
+        secondary_hue=gr.themes.colors.purple,
+        neutral_hue=gr.themes.colors.slate,
+        font=gr.themes.GoogleFont("Inter"),
+    ).set(
+        body_background_fill="#0f172a",
+        body_background_fill_dark="#0f172a",
+        block_background_fill="#1e293b",
+        block_background_fill_dark="#1e293b",
+        block_border_color="#334155",
+        block_border_color_dark="#334155",
+        block_label_text_color="#e2e8f0",
+        block_label_text_color_dark="#e2e8f0",
+        block_title_text_color="#e2e8f0",
+        block_title_text_color_dark="#e2e8f0",
+        input_background_fill="#0f172a",
+        input_background_fill_dark="#0f172a",
+        button_primary_background_fill="#3b82f6",
+        button_primary_background_fill_dark="#3b82f6",
+        button_primary_text_color="#ffffff",
+    )
+
+    with gr.Blocks(title="WhipStudio — ML Debug Arena") as app:
+
+        # ── Header ──
+        gr.HTML("""
+        <div class="arena-header">
+            <h1>🔧 WhipStudio — ML Debug Arena</h1>
+            <p>An RL environment where agents debug broken PyTorch training scripts</p>
+        </div>
+        """)
+
+        base_url = gr.Textbox(label="🌐 API Base URL", value=DEFAULT_BASE_URL, scale=1)
+
+        with gr.Row(equal_height=False):
+
+            # ── Left column: Task selector ──
+            with gr.Column(scale=1, min_width=280):
+                gr.Markdown("### 📋 Task Selector")
+                task_id = gr.Radio(
+                    choices=["task1", "task2", "task3", "task4", "task5"],
+                    value="task1",
+                    label="Select Task",
+                    info="Choose a debugging challenge",
+                )
+
+                task_desc = gr.Markdown(
+                    value="""### Broken Training Loop  🟢 Easy
+
+Fix optimizer order + learning rate bugs in a linear classifier.
+
+**💡 Hint:** Look at optimizer.step() / loss.backward() order and the learning rate."""
+                )
+
+                with gr.Row():
+                    btn_reset = gr.Button("🔄 Reset", variant="primary", size="sm")
+                    btn_baseline = gr.Button("🤖 Auto-Agent", variant="secondary", size="sm")
+
+                status = gr.Textbox(label="Status", interactive=False, lines=1)
+                timeline = gr.HTML(label="Episode Timeline", value=_step_timeline_html([], 0))
+
+            # ── Center column: Code editor ──
+            with gr.Column(scale=2, min_width=400):
+                gr.Markdown("### 💻 Code Editor")
+                code_editor = gr.Code(
+                    label="Your Fix (edit the code below)",
+                    language="python",
+                    lines=22,
+                )
+
+                with gr.Row():
+                    btn_submit = gr.Button("🚀 Submit Fix", variant="primary", size="lg")
+
+                error_log = gr.Textbox(
+                    label="📋 Execution Output / Error Log",
+                    lines=6,
+                    interactive=False,
+                )
+
+            # ── Right column: Results ──
+            with gr.Column(scale=1, min_width=300):
+                gr.Markdown("### 📊 Results")
+                score_display = gr.HTML(value=_score_html(0.0))
+
+                with gr.Tabs():
+                    with gr.Tab("📉 Loss Curve"):
+                        loss_plot = gr.Plot(label="Loss Curve")
+                    with gr.Tab("📈 Val Accuracy"):
+                        acc_plot = gr.Plot(label="Validation Accuracy (Task 3)")
+                    with gr.Tab("🔀 Code Diff"):
+                        diff_view = gr.HTML(
+                            value='<div style="color:#94a3b8;text-align:center;padding:20px;">Submit a fix to see the diff</div>'
+                        )
+                    with gr.Tab("🧭 Current State"):
+                        state_summary = gr.Markdown(
+                            value="Press Reset to start an episode, then Current State will appear here."
+                        )
+                        btn_refresh_state = gr.Button("🔄 Refresh State", variant="secondary", size="sm")
+                        state_json = gr.JSON(label="/state response", value={})
+
+                baseline_output = gr.Markdown(label="Baseline Results", visible=False)
+
+
+        # ── Bottom: Raw API tab (for developers) ──
+        with gr.Accordion("🔧 Developer Tools (Raw API)", open=False):
+            with gr.Row():
+                with gr.Column():
+                    dev_method = gr.Radio(["GET", "POST"], value="GET", label="Method")
+                    dev_path = gr.Textbox(label="Path", value="/health")
+                    dev_payload = gr.Code(label="Payload (JSON)", language="json", value="{}")
+                    btn_dev = gr.Button("Send Request", variant="secondary")
+                with gr.Column():
+                    dev_status = gr.Textbox(label="Status")
+                    dev_response = gr.Code(label="Response", language="json")
+
+            def dev_call(base, method, path, payload_text):
+                base = (base or DEFAULT_BASE_URL).strip().rstrip("/")
+                url = f"{base}{path}"
+                try:
+                    payload = json.loads(payload_text) if payload_text.strip() else {}
+                except json.JSONDecodeError as e:
+                    return f"JSON Error", str(e)
+                try:
+                    with httpx.Client(timeout=90.0) as client:
+                        if method == "GET":
+                            resp = client.get(url)
+                        else:
+                            resp = client.post(url, json=payload)
+                    ct = resp.headers.get("content-type", "")
+                    body = json.dumps(resp.json(), indent=2) if "json" in ct else resp.text
+                    return f"{resp.status_code} {resp.reason_phrase}", body
+                except Exception as exc:
+                    return "Error", f"{exc.__class__.__name__}: {exc}"
+
+            btn_dev.click(
+                fn=dev_call,
+                inputs=[base_url, dev_method, dev_path, dev_payload],
+                outputs=[dev_status, dev_response],
+            )
+
+
+        # ── Event bindings ──
+
+        # Task selector updates description
+        def update_task_desc(tid):
+            info = TASK_INFO.get(tid, {})
+            return f"""### {info.get('name', tid)}  {info.get('difficulty', '')}
+
+{info.get('description', '')}
+
+**💡 Hint:** {info.get('hints', 'No hints available.')}"""
+
+        task_id.change(fn=update_task_desc, inputs=[task_id], outputs=[task_desc])
+
+        # Reset
+        btn_reset.click(
+            fn=do_reset,
+            inputs=[base_url, task_id],
+            outputs=[status, code_editor, task_desc, score_display, loss_plot, acc_plot, diff_view, timeline, error_log],
+        ).then(
+            fn=load_current_state,
+            inputs=[base_url],
+            outputs=[state_summary, state_json],
+        )
+
+        # Submit fix
+        btn_submit.click(
+            fn=do_step,
+            inputs=[base_url, code_editor],
+            outputs=[status, score_display, loss_plot, acc_plot, diff_view, timeline, error_log],
+        ).then(
+            fn=load_current_state,
+            inputs=[base_url],
+            outputs=[state_summary, state_json],
+        )
+
+        btn_refresh_state.click(
+            fn=load_current_state,
+            inputs=[base_url],
+            outputs=[state_summary, state_json],
+        )
+
+        # Baseline (auto-agent) — live streaming per-task
+        TASK_NAMES = {
+            "task1": "🟢 Broken Training Loop",
+            "task2": "🟡 Silent NaN Loss",
+            "task3": "🔴 OOM + Data Leakage",
+            "task4": "🟡 Wrong Loss Function",
+            "task5": "🟡 Frozen Backbone",
+        }
+
+        def run_baseline_live(base_url_val):
+            """Generator that yields live progress as each task completes."""
+            base = (base_url_val or DEFAULT_BASE_URL).strip().rstrip("/")
+            results = {}
+            lines_header = ["### 🤖 Baseline Agent — Live Progress\n"]
+
+            # Phase 1: Show "starting" state
+            yield "\n".join(lines_header + ["⏳ Starting baseline agent..."])
+
+            for tid in ["task1", "task2", "task3", "task4", "task5"]:
+                tname = TASK_NAMES.get(tid, tid)
+
+                # Show "running this task" update
+                progress_lines = list(lines_header)
+                # Show completed tasks
+                for done_tid, info in results.items():
+                    s = info["score"]
+                    emoji = "🎯" if s >= 0.9 else ("✅" if s >= 0.7 else ("📈" if s >= 0.4 else "⚠️"))
+                    progress_lines.append(f"- {emoji} **{TASK_NAMES.get(done_tid, done_tid)}**: {s:.4f}")
+                    if info.get("error"):
+                        progress_lines.append(f"  - ⚠️ `{info['error'][:150]}`")
+                # Show currently running task
+                progress_lines.append(f"\n🤖 **Running {tname}** — agent is analyzing the code and generating a fix...")
+                progress_lines.append(f"\n*This may take 30-60 seconds per task (LLM inference + sandbox execution × 3 attempts)*")
+                yield "\n".join(progress_lines)
+
+                # Actually call the per-task endpoint
+                try:
+                    with httpx.Client(timeout=180.0) as client:
+                        resp = client.get(f"{base}/baseline/task/{tid}")
+                        resp.raise_for_status()
+                        data = resp.json()
+                except Exception as exc:
+                    data = {"score": 0.0, "error": f"{exc.__class__.__name__}: {exc}"}
+
+                results[tid] = {
+                    "score": float(data.get("score", 0.0)),
+                    "error": data.get("error", ""),
+                    "fixed_code": data.get("fixed_code", ""),
+                    "output": data.get("output", ""),
+                }
+
+            # Final summary
+            final_lines = ["### 🤖 Baseline Agent Results\n", "| Task | Score |", "|---|---|"]
+            total = 0.0
+            has_errors = False
+            for tid in ["task1", "task2", "task3", "task4", "task5"]:
+                info = results.get(tid, {"score": 0.0})
+                s = info["score"]
+                total += s
+                emoji = "🎯" if s >= 0.9 else ("✅" if s >= 0.7 else ("📈" if s >= 0.4 else "⚠️"))
+                final_lines.append(f"| {TASK_NAMES.get(tid, tid)} | {emoji} **{s:.4f}** |")
+                if info.get("error"):
+                    has_errors = True
+                    final_lines.append(f"\n> ⚠️ `{info['error'][:200]}`\n")
+
+            avg = total / 5
+            final_lines.append(f"\n**Average: {avg:.4f}**")
+            if avg >= 0.7:
+                final_lines.append("\n🎉 **Agent performed well!** The environment is solvable.")
+            elif avg >= 0.3:
+                final_lines.append("\n📈 **Agent showed partial progress.** Reward shaping is working.")
+            elif not has_errors:
+                final_lines.append("\n⚠️ **Agent scored low.** Tasks may be too challenging for zero-shot inference.")
+
+            if has_errors:
+                final_lines.append("\n---\n> [!WARNING]\n> Some tasks failed. Check if `HF_TOKEN` is valid and the model is accessible.")
+
+            final_lines.append("\n---\n### 🔍 Auto-Agent Generated Code & Execution Logs")
+            for tid in ["task1", "task2", "task3", "task4", "task5"]:
+                info = results.get(tid, {})
+                fixed_code = str(info.get("fixed_code", ""))
+                output = str(info.get("output", ""))
+                if fixed_code.strip() or output.strip():
+                    final_lines.append(f"\n#### {TASK_NAMES.get(tid, tid)}")
+                    if fixed_code.strip():
+                        final_lines.append("<details><summary><b>Show Generated Code</b></summary>\n\n```python\n" + fixed_code + "\n```\n</details>")
+                    if output.strip():
+                        final_lines.append("<details><summary><b>Show Execution Output</b></summary>\n\n```text\n" + output + "\n```\n</details>")
+
+            yield "\n".join(final_lines)
+
+        btn_baseline.click(
+            fn=lambda: gr.update(visible=True),
+            outputs=[baseline_output],
+        ).then(
+            fn=run_baseline_live,
+            inputs=[base_url],
+            outputs=[baseline_output],
+        )
+
+        # Footer
+        gr.HTML("""
+        <div style="text-align:center; padding:16px 0; color:#64748b; font-size:0.85em; border-top:1px solid #1e293b; margin-top:16px;">
+            WhipStudio v1.0 — OpenEnv ML Debug Environment
+            · <a href="/web" style="color:#60a5fa;">OpenEnv Web UI →</a>
+            · <a href="/docs" style="color:#60a5fa;">API Docs →</a>
+        </div>
+        """)
+
+    return app
+
+
+def main(host: str = "0.0.0.0", port: int = 7860):
+    app = build_ui()
+    app.launch(server_name=host, server_port=port, css=CUSTOM_CSS)
+
+
+if __name__ == "__main__":
+    main()

models.py

@@ -0,0 +1,41 @@
+from openenv.core.env_server.types import Action, Observation
+from pydantic import Field
+
+
+class MLDebugAction(Action):
+    """Agent submits a corrected training script."""
+
+    fixed_code: str = Field(
+        ...,
+        description="The corrected Python training script. Must be complete runnable code.",
+    )
+    explanation: str = Field(
+        default="",
+        description="Optional: agent's explanation of bugs found (not scored, for logging)",
+    )
+    attempt_number: int = Field(
+        default=1,
+        ge=1,
+        le=3,
+        description="Which attempt this is. Max 3 per episode.",
+    )
+
+
+class MLDebugObservation(Observation):
+    """What the agent sees on reset() and after each step()."""
+
+    task_id: str = Field(..., description="task1 | task2 | task3")
+    task_description: str = Field(..., description="Plain English task instructions")
+    buggy_code: str = Field(..., description="The broken training script")
+    error_log: str = Field(
+        default="",
+        description="stdout+stderr from the previous attempt. Empty on first step.",
+    )
+    last_reward: float = Field(
+        default=0.0,
+        description="Reward from previous attempt. 0.0 on first step.",
+    )
+    metrics: dict = Field(
+        default_factory=dict,
+        description="Structured: {final_loss, nan_count, val_acc, timed_out, exit_code}",
+    )

openenv.yaml

@@ -0,0 +1,77 @@
+
+spec_version: 2
+name: whipstudio-env
+type: space
+runtime: fastapi
+app: server.app:app
+port: 7860
+
+metadata:
+  title: WhipStudio Env
+  version: "1.1.0"
+  description: >
+    OpenEnv-compatible RL environment where agents debug broken PyTorch
+    training scripts across five tasks with continuous reward scoring.
+  author: Amogh-kal1
+  license: Apache-2.0
+  repository: Amogh-kal1/whipstudio-env
+  tags:
+    - openenv
+    - rl
+    - ml-debugging
+    - fastapi
+    - gradio
+
+compute:
+  python_version: "3.11"
+  cpu: 2
+  memory_gb: 16
+  gpu: false
+
+service:
+  host: 0.0.0.0
+  port: 7860
+  health_path: /health
+  base_path: /
+  web_path: /web
+  ui_path: /ui
+  docs_path: /docs
+  endpoints:
+    - method: POST
+      path: /reset
+    - method: POST
+      path: /step
+    - method: GET
+      path: /state
+    - method: GET
+      path: /tasks
+    - method: POST
+      path: /grader
+    - method: GET
+      path: /baseline
+    - method: GET
+      path: /baseline/task/{task_id}
+    - method: GET
+      path: /baseline/health
+
+tasks:
+  - id: task1
+    name: Broken training loop
+    difficulty: easy
+    max_steps: 3
+  - id: task2
+    name: Silent NaN loss
+    difficulty: medium
+    max_steps: 3
+  - id: task3
+    name: OOM and data leakage
+    difficulty: hard
+    max_steps: 3
+  - id: task4
+    name: Wrong loss function
+    difficulty: medium
+    max_steps: 3
+  - id: task5
+    name: Frozen backbone
+    difficulty: medium
+    max_steps: 3

openenv_whipstudio.egg-info/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.4
+Name: openenv-whipstudio
+Version: 0.1.0
+Summary: ML Debug environment for OpenEnv
+Requires-Python: >=3.11
+License-File: LICENSE
+Requires-Dist: openenv-core[core]>=0.2.1
+Requires-Dist: fastapi>=0.110.0
+Requires-Dist: uvicorn>=0.27.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: httpx>=0.27.0
+Requires-Dist: torch>=2.2.0
+Requires-Dist: smolagents>=1.0.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Dynamic: license-file

openenv_whipstudio.egg-info/SOURCES.txt

@@ -0,0 +1,25 @@
+LICENSE
+README.md
+__init__.py
+baseline_agent.py
+client.py
+gradio_app.py
+models.py
+pyproject.toml
+openenv_whipstudio.egg-info/PKG-INFO
+openenv_whipstudio.egg-info/SOURCES.txt
+openenv_whipstudio.egg-info/dependency_links.txt
+openenv_whipstudio.egg-info/entry_points.txt
+openenv_whipstudio.egg-info/requires.txt
+openenv_whipstudio.egg-info/top_level.txt
+server/__init__.py
+server/app.py
+server/environment.py
+server/sandbox.py
+server/tasks/__init__.py
+server/tasks/graders.py
+server/tasks/task1_broken_loop.py
+server/tasks/task2_nan_loss.py
+server/tasks/task3_oom_leakage.py
+server/tasks/task4_wrong_loss.py
+server/tasks/task5_frozen_backbone.py

openenv_whipstudio.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

openenv_whipstudio.egg-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+server = server.app:main
+whipstudio-server = server.app:main

openenv_whipstudio.egg-info/requires.txt

@@ -0,0 +1,10 @@
+openenv-core[core]>=0.2.1
+fastapi>=0.110.0
+uvicorn>=0.27.0
+pydantic>=2.0.0
+httpx>=0.27.0
+torch>=2.2.0
+smolagents>=1.0.0
+
+[dev]
+pytest>=8.0.0

openenv_whipstudio.egg-info/top_level.txt

@@ -0,0 +1 @@
+server

project_status_report_2703.md

@@ -0,0 +1,70 @@
+# WhipStudio Project Status Report
+
+**Date:** March 27, 2026  
+**Project:** WhipStudio — ML Debugging Arena (OpenEnv AI Hackathon)  
+**Status:** 🟢 **On Track for Submission**
+
+---
+
+## Executive Summary
+WhipStudio has successfully transitioned from a conceptual prototype to a fully functional, professional-grade ML Debugging Arena. The core objective of the platform—providing an isolated, standardized environment (Gymnasium) to evaluate and train autonomous AI agents on PyTorch debugging tasks—is now fully operational. The backend conforms exactly to the OpenEnv specification, while the frontend provides a best-in-class spectator and verification experience.
+
+---
+
+## System Architecture & Technical Specifications
+
+### 1. The Core Environment Database (Tasks)
+The platform currently ships with three distinct ML debugging challenges, graded automatically:
+- **Task 1 (Easy): Broken Training Loop.** Fixes a simple linear classifier that has an excessively high learning rate, steps the optimizer out of order, and computes loss incorrectly.
+- **Task 2 (Medium): Silent NaN Loss.** Fixes a CNN training script where the output loss silently turns to `NaN`.
+- **Task 3 (Hard): OOM + Data Leakage.** Fixes a complex PyTorch training loop that creates memory leaks by accumulating computation graphs incorrectly, and introduces data leakage by augmenting before splitting train/val. Graded on a granular scale (0.15 - 0.5 per bug fixed).
+
+### 2. The Execution Sandbox API (Backend)
+All agent submissions are processed securely in an isolated environment.
+
+**Core Endpoints (OpenEnv Core API):**
+- `POST /reset`: Initialize the environment for a specific task. Returns the `observation` payload, which includes the `buggy_code` string and `task_description`.
+- `POST /step`: Submit a fixed code attempt (the `action`). The system executes the code in the sandbox, parses the `stdout`/`stderr` using regex metrics extractors (e.g., [parse_scalar](file:///home/amogh/Documents/openenv-comp/WhipStudio/server/tasks/graders.py#35-38) for `FINAL_LOSS`), scores the execution, logs the trajectory, and returns the formal RL `reward` (0.0 to 1.0) along with the `done` state.
+
+**Security & Isolation ([server/sandbox.py](file:///home/amogh/Documents/openenv-comp/WhipStudio/server/sandbox.py)):**
+- Python execution restricts dangerous operations using `BANNED_PATTERNS` (e.g., `os.system`, `subprocess.`, `open()`, `socket.`, `requests.`).
+- Code runs in `/tmp` natively in a subprocess with a strict 30-second timeout (`TIMEOUT_SECONDS = 30`) and output sizing limits (`MAX_OUTPUT_BYTES = 8192`). 
+- Environment dependencies (`PYTHONPATH`) are strictly controlled, ensuring library access (like `torch`) without host system compromise.
+
+### 3. AI Agent Integrations
+- `GET /baseline/health`: Checks if the default baseline model (Qwen-32B) is accessible and authenticated via the HuggingFace API using the `.env` configuration.
+- `GET /baseline`: Instructs the backend to run the internal autonomous zero-shot agent across all three tasks simultaneously.
+- `GET /baseline/task/{task_id}`: Streams execution of the baseline agent for a specific task (120s timeout). Used extensively by the Gradio UI for real-time visualization.
+
+### 4. User Interfaces (Frontend)
+
+WhipStudio provides dual interfaces for varying stakeholder needs:
+
+**A. The "ML Debug Arena" Verification Portal (Gradio App)**
+A custom-built, interactive dashboard designed for human verification and hackathon judges. 
+- **Code Editor:** Monaco-style interactive Python editor pre-filled with the buggy code payload.
+- **Live Loss Visualization:** Integrates `matplotlib` to plot `LOSS` curves and `VAL_ACCURACY` step-by-step from the agent's debugged submission.
+- **Live Diff Context:** A source-control style diff view showing exactly what code lines the agent inserted or deleted to achieve the fix.
+- **Autonomous Streaming Mode:** The "Auto-Agent" trigger connects to the `/baseline/task/{task_id}` streams to provide a live "Chain-of-Thought" experience, displaying progress and rewards directly to the UI panel in real-time.
+
+**B. The OpenEnv Compliance Web UI (`/web`)**
+- Directly embedded via the OpenEnv python framework (`os.environ["ENABLE_WEB_INTERFACE"] = "true"`).
+- Accessed via `http://[host]:8000/web`, this provides the raw, standardized OpenEnv chat-style UI validating 100% adherence to the hackathon's core platform requirements.
+
+---
+
+## Completed Milestones
+
+1. **Bug Remediation:** Fixed critical parsing logic failures where double-escaped regex patterns (`[\\\\d\\\\.\\\\-]+`) prevented the extraction of metrics like `FINAL_LOSS:0.4523`. Also resolved the `ModuleNotFoundError: torch` in the sandbox.
+2. **Environment Token Logic:** Re-wrote authentication management using `python-dotenv` for local caching overrides, ensuring 401 Unauthorized API failures do not disrupt model inference when tokens expire mid-session.
+3. **Task Trajectory Generation:** Enabled persistent trajectory state logging inside `environment.py` for advanced analytics and training extraction.
+
+---
+
+## Strategic Next Steps
+
+**Phase 2: Reinforcement Learning Integration**
+The zero-shot inference pipeline currently scores dynamically. The immediate next phase leverages the Trajectory Logic:
+1. Export the stored trajectory JSON logs containing `buggy_code`, `reward`, and `error_log`.
+2. Integrate the **TRL (Transformer Reinforcement Learning) library**.
+3. Implement Group Relative Policy Optimization (GRPO) to iteratively train a specific lightweight model (like a 7B parameter local model) to perform significantly better on PyTorch debugging tasks than larger, generalized zero-shot models.

pyproject.toml

@@ -0,0 +1,31 @@
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-whipstudio"
+version = "0.1.0"
+description = "ML Debug environment for OpenEnv"
+requires-python = ">=3.11"
+dependencies = [
+    "openenv-core[core]>=0.2.1",
+    "fastapi>=0.110.0",
+    "uvicorn>=0.27.0",
+    "pydantic>=2.0.0",
+    "httpx>=0.27.0",
+    "torch>=2.2.0",
+    "smolagents>=1.0.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+]
+
+[project.scripts]
+whipstudio-server = "server.app:main"
+server = "server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["server", "server.tasks"]

server/__init__.py

@@ -0,0 +1 @@
+"""ML debug server package."""

server/app.py

@@ -0,0 +1,254 @@
+import asyncio
+import os
+import sys
+
+from dotenv import load_dotenv
+load_dotenv(override=True)
+
+import httpx
+from fastapi import FastAPI, Request
+from fastapi.responses import RedirectResponse
+from fastapi.responses import HTMLResponse
+
+from openenv.core.env_server.http_server import create_app
+
+_project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if _project_root not in sys.path:
+    sys.path.insert(0, _project_root)
+
+try:
+    from ..models import MLDebugAction, MLDebugObservation
+    from .environment import MLDebugEnvironment
+    from .tasks.graders import RunResult, score_task
+except ImportError:
+    from models import MLDebugAction, MLDebugObservation
+    from server.environment import MLDebugEnvironment
+    from server.tasks.graders import RunResult, score_task
+
+# Disable OpenEnv's default web UI so /web can mirror the custom Gradio UI.
+os.environ["ENABLE_WEB_INTERFACE"] = "false"
+
+app: FastAPI = create_app(
+    MLDebugEnvironment,
+    MLDebugAction,
+    MLDebugObservation,
+    env_name="whipstudio-env",
+    max_concurrent_envs=4,
+)
+
+
+@app.get("/__build", include_in_schema=False)
+def build_info():
+    """Build/runtime fingerprint to confirm what code is deployed."""
+    import platform
+
+    return {
+        "env_name": "whipstudio-env",
+        "python": platform.python_version(),
+        "platform": platform.platform(),
+        "port": os.environ.get("PORT"),
+        "enable_web_interface": os.environ.get("ENABLE_WEB_INTERFACE"),
+    }
+
+
+def _has_route(path: str, method: str) -> bool:
+    method = method.upper()
+    for route in app.router.routes:
+        if getattr(route, "path", None) != path:
+            continue
+        methods = getattr(route, "methods", None)
+        if methods and method in methods:
+            return True
+    return False
+
+
+@app.get("/", include_in_schema=False)
+def root_redirect():
+    return RedirectResponse(url="/ui", status_code=307)
+
+
+if not _has_route("/health", "GET"):
+
+    @app.get("/health", include_in_schema=False)
+    def health_get():
+        return {"status": "ok"}
+
+
+if not _has_route("/health", "POST"):
+
+    @app.post("/health", include_in_schema=False)
+    def health_post():
+        return {"status": "ok"}
+
+
+@app.get("/reset")
+def reset_liveness():
+    return {"status": "ok", "message": "use POST /reset to start an episode"}
+
+
+@app.get("/tasks")
+def list_tasks():
+    return {
+        "tasks": [
+            {"id": "task1", "name": "Broken training loop", "difficulty": "easy"},
+            {"id": "task2", "name": "Silent NaN loss", "difficulty": "medium"},
+            {"id": "task3", "name": "OOM and data leakage", "difficulty": "hard"},
+            {"id": "task4", "name": "Wrong loss function", "difficulty": "medium"},
+            {"id": "task5", "name": "Frozen backbone", "difficulty": "medium"},
+        ],
+        "action_schema": {
+            "fixed_code": "string (required) — complete runnable Python script",
+            "explanation": "string (optional) — description of bugs found",
+            "attempt_number": "int 1-3 (optional) — which attempt this is",
+        },
+    }
+
+
+@app.post("/grader")
+def run_grader(payload: dict):
+    task_id = payload.get("task_id", "task1")
+    result = RunResult(
+        exit_code=payload.get("exit_code", -1),
+        stdout=payload.get("stdout", ""),
+        stderr=payload.get("stderr", ""),
+        elapsed_seconds=payload.get("elapsed", 0.0),
+        timed_out=payload.get("timed_out", False),
+        fixed_code=payload.get("fixed_code", ""),
+    )
+    score, breakdown = score_task(task_id, result)
+    return {"task_id": task_id, "score": score, "breakdown": breakdown}
+
+
+@app.get("/baseline")
+async def run_baseline(request: Request):
+    try:
+        from ..baseline_agent import run_single_task
+    except ImportError:
+        from baseline_agent import run_single_task
+
+    env_url = str(request.base_url).rstrip("/")
+    results = {}
+    task_scores = {}
+    for task_id in ["task1", "task2", "task3", "task4", "task5"]:
+        try:
+            score = await asyncio.wait_for(run_single_task(task_id, env_url), timeout=120.0)
+            results[task_id] = round(score, 4)
+            task_scores[task_id] = round(score, 4)
+        except TimeoutError:
+            results[task_id] = 0.0
+            task_scores[task_id] = 0.0
+            results[f"{task_id}_error"] = "timeout: task took longer than 120s"
+        except httpx.HTTPError as exc:
+            results[task_id] = 0.0
+            task_scores[task_id] = 0.0
+            results[f"{task_id}_error"] = f"http_error: {exc.__class__.__name__}: {exc}"
+        except Exception as exc:
+            results[task_id] = 0.0
+            task_scores[task_id] = 0.0
+            results[f"{task_id}_error"] = f"internal_error: {exc.__class__.__name__}: {exc}"
+    avg = round(sum(task_scores.values()) / max(1, len(task_scores)), 4)
+    return {"baseline_scores": results, "average": avg, "env_url": env_url}
+
+
+@app.get("/baseline/task/{task_id}")
+async def run_baseline_single(task_id: str, request: Request):
+    """Run the baseline agent on a single task. Returns score + details."""
+    try:
+        from ..baseline_agent import run_single_task_detailed
+    except ImportError:
+        from baseline_agent import run_single_task_detailed
+
+    env_url = str(request.base_url).rstrip("/")
+    try:
+        result = await asyncio.wait_for(run_single_task_detailed(task_id, env_url), timeout=120.0)
+        return {
+            "task_id": task_id,
+            "score": round(result["score"], 4),
+            "status": "ok",
+            "fixed_code": result.get("fixed_code", ""),
+            "output": result.get("output", ""),
+            "attempts": result.get("attempts", []),
+        }
+    except TimeoutError:
+        return {"task_id": task_id, "score": 0.0, "status": "timeout", "error": "Task took longer than 120s"}
+    except Exception as exc:
+        return {"task_id": task_id, "score": 0.0, "status": "error", "error": f"{exc.__class__.__name__}: {exc}"}
+
+
+@app.get("/baseline/health")
+def baseline_health():
+    hf_token_present = bool(os.environ.get("HF_TOKEN"))
+    model_ready = False
+    model_error = None
+
+    try:
+        try:
+            from ..baseline_agent import get_model
+        except ImportError:
+            from baseline_agent import get_model
+
+        get_model()
+        model_ready = True
+    except Exception as exc:
+        model_error = f"{exc.__class__.__name__}: {exc}"
+
+    status = "ok" if hf_token_present and model_ready else "degraded"
+    return {
+        "status": status,
+        "hf_token_present": hf_token_present,
+        "model_ready": model_ready,
+        "model_error": model_error,
+    }
+
+
+_ui_mounted = False
+try:
+    import gradio as gr
+    try:
+        from ..gradio_app import build_ui
+    except ImportError:
+        from gradio_app import build_ui
+
+    gradio_ui = build_ui()
+    app = gr.mount_gradio_app(app, gradio_ui, path="/ui")
+    _ui_mounted = True
+except Exception as e:
+    # Don't fail silently in Spaces: return a helpful error page at /ui.
+    import traceback
+
+    print(f"Failed to mount Gradio UI: {e}")
+    traceback.print_exc()
+
+
+if not _ui_mounted:
+    @app.get("/ui", include_in_schema=False)
+    @app.get("/ui/", include_in_schema=False)
+    def ui_mount_failed():
+        return HTMLResponse(
+            "<h2>WhipStudio UI failed to start</h2>"
+            "<p>The API server is running, but the Gradio UI could not be mounted.</p>"
+            "<p>Check container logs for <code>Failed to mount Gradio UI</code>.</p>",
+            status_code=500,
+        )
+
+
+@app.api_route("/web", methods=["GET", "POST"], include_in_schema=False)
+def web_redirect_root():
+    return RedirectResponse(url="/ui", status_code=307)
+
+
+@app.api_route("/web/{path:path}", methods=["GET", "POST"], include_in_schema=False)
+def web_redirect_path(path: str):
+    if path:
+        return RedirectResponse(url=f"/ui/{path}", status_code=307)
+    return RedirectResponse(url="/ui", status_code=307)
+
+
+def main(host: str = "0.0.0.0", port: int = 7860):
+    import uvicorn
+
+    uvicorn.run("server.app:app", host=host, port=port, reload=False)
+
+
+if __name__ == "__main__":
+    main()

server/environment.py

@@ -0,0 +1,165 @@
+import re
+import math
+import time
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+try:
+    from ..models import MLDebugAction, MLDebugObservation
+    from .sandbox import execute_code
+    from .tasks import task1_broken_loop, task2_nan_loss, task3_oom_leakage, task4_wrong_loss, task5_frozen_backbone
+    from .tasks.graders import parse_losses, parse_val_accs, score_task
+except ImportError:
+    from models import MLDebugAction, MLDebugObservation
+    from server.sandbox import execute_code
+    from server.tasks import task1_broken_loop, task2_nan_loss, task3_oom_leakage, task4_wrong_loss, task5_frozen_backbone
+    from server.tasks.graders import parse_losses, parse_val_accs, score_task
+
+TASKS = {
+    "task1": task1_broken_loop,
+    "task2": task2_nan_loss,
+    "task3": task3_oom_leakage,
+    "task4": task4_wrong_loss,
+    "task5": task5_frozen_backbone,
+}
+
+
+class MLDebugEnvironment(Environment):
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._task_id = "task1"
+        self._best_reward = 0.0
+        self._trajectory: list[dict] = []
+
+    def reset(self, task_id: str = "task1", **kwargs) -> MLDebugObservation:  # type: ignore[override]
+        if task_id not in TASKS:
+            task_id = "task1"
+
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._task_id = task_id
+        self._best_reward = 0.0
+        self._trajectory = []
+
+        task = TASKS[self._task_id]
+        return MLDebugObservation(
+            task_id=self._task_id,
+            task_description=task.TASK_DESCRIPTION.strip(),
+            buggy_code=task.BUGGY_CODE.strip(),
+            error_log="",
+            last_reward=0.0,
+            metrics={},
+            done=False,
+            reward=0.0,
+        )
+
+    def step(self, action: MLDebugAction) -> MLDebugObservation:  # type: ignore[override]
+        self._state.step_count += 1
+
+        if not action.fixed_code or not action.fixed_code.strip():
+            done = True
+            metrics = {
+                "exit_code": -1,
+                "elapsed_seconds": 0.0,
+                "timed_out": False,
+                "step": self._state.step_count,
+                "best_reward_so_far": self._best_reward,
+                "error": "empty code submitted",
+            }
+            task = TASKS[self._task_id]
+            return MLDebugObservation(
+                task_id=self._task_id,
+                task_description=task.TASK_DESCRIPTION.strip(),
+                buggy_code=task.BUGGY_CODE.strip(),
+                error_log="empty code submitted",
+                last_reward=0.0,
+                metrics=metrics,
+                done=done,
+                reward=0.0,
+            )
+
+        run_result1 = execute_code(action.fixed_code)
+        reward1, breakdown1 = score_task(self._task_id, run_result1)
+
+        consistency_flag = False
+        reward_variance = 0.0
+        final_reward = reward1
+        final_breakdown = breakdown1
+        run_result = run_result1
+
+        if reward1 > 0.5:
+            run_result2 = execute_code(action.fixed_code)
+            reward2, breakdown2 = score_task(self._task_id, run_result2)
+            reward_variance = abs(reward1 - reward2)
+            if reward_variance > 0.15:
+                consistency_flag = True
+                final_reward = min(reward1, reward2)
+                if reward2 < reward1:
+                    final_breakdown = breakdown2
+                    run_result = run_result2
+            else:
+                consistency_flag = False
+                final_reward = (reward1 + reward2) / 2.0
+        else:
+            consistency_flag = False
+            final_reward = reward1
+
+        self._best_reward = max(self._best_reward, final_reward)
+        done = self._state.step_count >= 3 or final_reward >= 0.95
+
+        losses = parse_losses(run_result.stdout)
+        val_accs = parse_val_accs(run_result.stdout)
+        final_loss = None
+        if losses:
+            final_loss = losses[-1]
+        else:
+            match = re.search(r"FINAL_LOSS:([-\d.]+)", run_result.stdout)
+            if match:
+                final_loss = float(match.group(1))
+
+        metrics = {
+            "exit_code": run_result.exit_code,
+            "elapsed_seconds": run_result.elapsed_seconds,
+            "timed_out": run_result.timed_out,
+            "step": self._state.step_count,
+            "best_reward_so_far": self._best_reward,
+            "final_loss": final_loss,
+            "nan_count": sum(1 for x in losses if math.isnan(x) or math.isinf(x)) if losses else 0,
+            "val_acc": val_accs[-1] if val_accs else None,
+            "consistency_flag": consistency_flag,
+            "reward_variance": round(reward_variance, 4),
+            "reward_breakdown": final_breakdown,
+        }
+
+        task = TASKS[self._task_id]
+
+        self._trajectory.append({
+            "step": self._state.step_count,
+            "reward": final_reward,
+            "best_reward": self._best_reward,
+            "metrics": metrics,
+            "done": done,
+            "timestamp": time.time(),
+        })
+
+        return MLDebugObservation(
+            task_id=self._task_id,
+            task_description=task.TASK_DESCRIPTION.strip(),
+            buggy_code=task.BUGGY_CODE.strip(),
+            error_log=(run_result.stdout + "\n" + run_result.stderr).strip()[:2000],
+            last_reward=final_reward,
+            metrics=metrics,
+            done=done,
+            reward=final_reward,
+        )
+
+    @property
+    def trajectory(self) -> list[dict]:
+        return list(self._trajectory)
+
+    @property
+    def state(self) -> State:
+        return self._state

server/requirements.txt

@@ -0,0 +1,11 @@
+openenv-core>=0.1.1
+fastapi>=0.110.0
+uvicorn>=0.27.0
+httpx>=0.27.0
+smolagents>=1.0.0
+pydantic>=2.0.0
+python-dotenv>=1.0.0
+gradio>=4.0.0
+tqdm>=4.0.0
+scikit-learn
+matplotlib

server/sandbox.py

@@ -0,0 +1,88 @@
+import os
+import subprocess
+import sys
+import tempfile
+import time
+
+from .tasks.graders import RunResult
+
+TIMEOUT_SECONDS = 30
+MAX_OUTPUT_BYTES = 8192
+BANNED_PATTERNS = [
+    "os.system",
+    "subprocess.",
+    "shutil.rmtree",
+    "open(",
+    "__import__",
+    "exec(",
+    "socket.",
+    "urllib.",
+    "requests.",
+]
+SAFE_ENV = {
+    "PATH": os.environ.get("PATH", "/usr/bin:/usr/local/bin"),
+    "HOME": "/tmp",
+    "PYTHONPATH": os.pathsep.join(sys.path),
+    "PYTHONDONTWRITEBYTECODE": "1",
+}
+
+
+def strip_markdown_code(code: str) -> str:
+    if "```python" in code:
+        return code.split("```python", 1)[1].split("```", 1)[0].strip()
+    if "```" in code:
+        return code.split("```", 1)[1].split("```", 1)[0].strip()
+    return code.strip()
+
+
+def execute_code(code: str) -> RunResult:
+    """Execute agent-submitted code in an isolated subprocess."""
+    cleaned_code = strip_markdown_code(code)
+
+    for pattern in BANNED_PATTERNS:
+        if pattern in cleaned_code:
+            return RunResult(
+                exit_code=-1,
+                stdout="",
+                stderr=f'Execution blocked: banned pattern "{pattern}" detected.',
+                elapsed_seconds=0.0,
+                timed_out=False,
+                fixed_code=cleaned_code,
+            )
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, dir="/tmp") as temp_file:
+        temp_file.write(cleaned_code)
+        tmp_path = temp_file.name
+
+    start = time.time()
+    try:
+        proc = subprocess.run(
+            ["python", tmp_path],
+            capture_output=True,
+            text=True,
+            timeout=TIMEOUT_SECONDS,
+            env=SAFE_ENV,
+            cwd="/tmp",
+        )
+        return RunResult(
+            exit_code=proc.returncode,
+            stdout=proc.stdout[:MAX_OUTPUT_BYTES],
+            stderr=proc.stderr[:2048],
+            elapsed_seconds=round(time.time() - start, 2),
+            timed_out=False,
+            fixed_code=cleaned_code,
+        )
+    except subprocess.TimeoutExpired:
+        return RunResult(
+            exit_code=-1,
+            stdout="",
+            stderr="Execution timed out after 30 seconds.",
+            elapsed_seconds=TIMEOUT_SECONDS,
+            timed_out=True,
+            fixed_code=cleaned_code,
+        )
+    finally:
+        try:
+            os.unlink(tmp_path)
+        except Exception:
+            pass

server/start.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+set -e
+
+export ENABLE_WEB_INTERFACE=false
+
+# Single process: FastAPI serves both API and mounted Gradio at /ui
+uvicorn server.app:app --host 0.0.0.0 --port "${PORT:-7860}"

server/tasks/__init__.py

@@ -0,0 +1 @@
+"""Task definitions for ML debug environment."""

server/tasks/graders.py

@@ -0,0 +1,256 @@
+import math
+import re
+import ast
+from dataclasses import dataclass
+
+
+@dataclass
+class RunResult:
+    exit_code: int
+    stdout: str
+    stderr: str
+    elapsed_seconds: float
+    timed_out: bool
+    fixed_code: str = ""
+
+
+def extract_metrics_block(stdout: str) -> str:
+    match = re.search(r"##METRICS_START##(.*?)##METRICS_END##", stdout, re.DOTALL)
+    if match:
+        return match.group(1)
+    return stdout
+
+
+def parse_losses(stdout: str) -> list[float]:
+    stdout = extract_metrics_block(stdout)
+    match = re.search(r"LOSSES:\[([^\]]+)\]", stdout)
+    if not match:
+        return []
+    try:
+        return [float(x.strip()) for x in match.group(1).split(",")]
+    except Exception:
+        return []
+
+
+def parse_val_accs(stdout: str) -> list[float]:
+    stdout = extract_metrics_block(stdout)
+    match = re.search(r"VAL_ACCS:\[([^\]]+)\]", stdout)
+    if not match:
+        return []
+    try:
+        return [float(x.strip()) for x in match.group(1).split(",")]
+    except Exception:
+        return []
+
+
+def parse_scalar(stdout: str, key: str) -> float | None:
+    stdout = extract_metrics_block(stdout)
+    match = re.search(rf"{key}:([-\d.]+)", stdout)
+    return float(match.group(1)) if match else None
+
+
+def is_valid_submission(code: str, stdout: str, exit_code: int) -> tuple[bool, str]:
+    if exit_code == 0:
+        if "LOSSES:" not in stdout and "FINAL_LOSS:" not in stdout:
+            return False, "No valid metrics output detected"
+        if "LOSSES:" in stdout:
+            losses = parse_losses(stdout)
+            if len(losses) < 5:
+                return False, "Fewer than 5 loss values parsed"
+    try:
+        tree = ast.parse(code)
+        if not any(isinstance(node, (ast.For, ast.While)) for node in ast.walk(tree)):
+            return False, "No ast.For or ast.While node found"
+    except Exception:
+        pass
+    return True, ""
+
+
+def sigmoid_reward(value: float, center: float, steepness: float, invert: bool = False) -> float:
+    try:
+        if invert:
+            x = steepness * (value - center)
+        else:
+            x = steepness * (center - value)
+        return round(1.0 / (1.0 + math.exp(-x)), 4)
+    except OverflowError:
+        return 0.0 if (invert and value > center) or (not invert and value < center) else 1.0
+
+
+def grade_task1(result: RunResult) -> tuple[float, dict]:
+    valid, reason = is_valid_submission(result.fixed_code, result.stdout, result.exit_code)
+    if not valid:
+        return 0.0, {"reason": reason}
+
+    if result.timed_out:
+        return 0.05, {"reason": "timed_out"}
+    if result.exit_code != 0:
+        return 0.0, {"reason": "crash"}
+
+    losses = parse_losses(result.stdout)
+    if not losses:
+        return 0.1, {"reason": "no_losses_parsed"}
+    if any(math.isnan(loss) or math.isinf(loss) for loss in losses):
+        return 0.15, {"reason": "nan_inf_found"}
+
+    final = losses[-1]
+    base_score = sigmoid_reward(final, center=0.75, steepness=3.0, invert=True)
+
+    bonus = 0.0
+    half = len(losses) // 2
+    if half > 0:
+        first_half = sum(losses[:half]) / half
+        second_half = sum(losses[half:]) / len(losses[half:])
+        if second_half < 0.85 * first_half:
+            bonus = 0.1
+            
+    final_score = min(1.0, base_score + bonus)
+    breakdown = {"base_score": base_score, "monotonicity_bonus": bonus}
+    return final_score, breakdown
+
+
+def grade_task2(result: RunResult) -> tuple[float, dict]:
+    valid, reason = is_valid_submission(result.fixed_code, result.stdout, result.exit_code)
+    if not valid:
+        return 0.0, {"reason": reason}
+
+    if result.timed_out:
+        return 0.05, {"reason": "timed_out"}
+    if result.exit_code != 0:
+        return 0.0, {"reason": "crash"}
+
+    losses = parse_losses(result.stdout)
+    if not losses or len(losses) < 30:
+        return 0.1, {"reason": "too_few_losses"}
+
+    nan_count = sum(1 for loss in losses if math.isnan(loss) or math.isinf(loss))
+    if nan_count == len(losses):
+        return 0.0, {"reason": "all_nans"}
+        
+    nan_ratio = nan_count / len(losses)
+    finite_losses = [loss for loss in losses if not math.isnan(loss) and not math.isinf(loss)]
+    final_finite_loss = finite_losses[-1] if finite_losses else float('inf')
+
+    convergence_score = sigmoid_reward(final_finite_loss, center=0.5, steepness=4.0, invert=True)
+    convergence_score *= (1.0 - nan_ratio)
+
+    stability_bonus = 0.0
+    if len(finite_losses) >= 20:
+        tail = finite_losses[-20:]
+        mean_tail = sum(tail) / len(tail)
+        tail_variance = sum((x - mean_tail) ** 2 for x in tail) / len(tail)
+        stability_bonus = sigmoid_reward(tail_variance, center=0.01, steepness=200.0, invert=True) * 0.1
+
+    final_score = min(1.0, convergence_score + stability_bonus)
+    breakdown = {"convergence_score": convergence_score, "nan_penalty": (1.0 - nan_ratio), "stability_bonus": stability_bonus, "nan_ratio": nan_ratio}
+    return final_score, breakdown
+
+
+def grade_task3(result: RunResult) -> tuple[float, dict]:
+    valid, reason = is_valid_submission(result.fixed_code, result.stdout, result.exit_code)
+    if not valid:
+        return 0.0, {"reason": reason}
+
+    if result.timed_out:
+        return 0.1, {"reason": "timed_out"}
+
+    if result.exit_code != 0:
+        if "out of memory" in result.stderr.lower():
+            return 0.1, {"reason": "oom"}
+        return 0.0, {"reason": "crash"}
+
+    val_accs = parse_val_accs(result.stdout)
+    final_loss_val = parse_scalar(result.stdout, "FINAL_LOSS")
+
+    memory_score = 0.0
+    if final_loss_val is not None:
+        memory_score = sigmoid_reward(final_loss_val, center=50.0, steepness=0.05, invert=True) * 0.5
+
+    leakage_score = 0.0
+    early_acc = 0.0
+    final_acc = 0.0
+    if val_accs and len(val_accs) >= 2:
+        early_acc = sum(val_accs[:2]) / 2.0
+        final_acc = val_accs[-1]
+        
+        leak_p1 = sigmoid_reward(early_acc, center=0.75, steepness=20.0, invert=True) * 0.3
+        leak_p2 = sigmoid_reward(final_acc, center=0.68, steepness=15.0, invert=False) * 0.7
+        leakage_score = (leak_p1 + leak_p2) * 0.5
+
+    final_score = min(1.0, memory_score + leakage_score)
+    breakdown = {"memory_score": memory_score, "leakage_score": leakage_score, "early_acc": early_acc, "final_acc": final_acc}
+    return final_score, breakdown
+
+
+def grade_task4(result: RunResult) -> tuple[float, dict]:
+    valid, reason = is_valid_submission(result.fixed_code, result.stdout, result.exit_code)
+    if not valid:
+        return 0.0, {"reason": reason}
+
+    if result.timed_out:
+        return 0.1, {"reason": "timed_out"}
+
+    if result.exit_code != 0:
+        return 0.0, {"reason": "crash"}
+
+    final_loss = parse_scalar(result.stdout, "FINAL_LOSS")
+    avg_labels = parse_scalar(result.stdout, "AVG_LABELS")
+    f1 = parse_scalar(result.stdout, "F1_SCORE")
+
+    loss_score = 0.0
+    if final_loss is not None:
+        loss_score = sigmoid_reward(final_loss, center=0.5, steepness=4.0, invert=True) * 0.3
+
+    labels_score = 0.0
+    if avg_labels is not None:
+        labels_score = sigmoid_reward(avg_labels, center=1.0, steepness=5.0, invert=False) * 0.3
+
+    f1_s = 0.0
+    if f1 is not None:
+        f1_s = sigmoid_reward(f1, center=0.6, steepness=10.0, invert=False) * 0.4
+
+    final_score = min(1.0, loss_score + labels_score + f1_s)
+    breakdown = {"loss_score": loss_score, "labels_score": labels_score, "f1_score": f1_s}
+    return final_score, breakdown
+
+
+def grade_task5(result: RunResult) -> tuple[float, dict]:
+    valid, reason = is_valid_submission(result.fixed_code, result.stdout, result.exit_code)
+    if not valid:
+        return 0.0, {"reason": reason}
+
+    if result.timed_out:
+        return 0.1, {"reason": "timed_out"}
+
+    if result.exit_code != 0:
+        return 0.0, {"reason": "crash"}
+
+    final_loss = parse_scalar(result.stdout, "FINAL_LOSS")
+    grad_norm = parse_scalar(result.stdout, "BACKBONE_GRAD_NORM")
+
+    loss_score = 0.0
+    if final_loss is not None:
+        loss_score = sigmoid_reward(final_loss, center=2.2, steepness=3.0, invert=True) * 0.5
+        
+    grad_score = 0.0
+    if grad_norm is not None:
+        grad_score = sigmoid_reward(grad_norm, center=0.001, steepness=1000.0, invert=False) * 0.5
+
+    final_score = min(1.0, loss_score + grad_score)
+    breakdown = {"loss_score": loss_score, "grad_score": grad_score}
+    return final_score, breakdown
+
+
+def score_task(task_id: str, result: RunResult) -> tuple[float, dict]:
+    graders = {
+        "task1": grade_task1,
+        "task2": grade_task2,
+        "task3": grade_task3,
+        "task4": grade_task4,
+        "task5": grade_task5,
+    }
+    if task_id not in graders:
+        raise ValueError(f"Unknown task_id: {task_id}")
+
+    score, breakdown = graders[task_id](result)
+    return round(max(0.0, min(1.0, score)), 4), breakdown

server/tasks/task1_broken_loop.py

@@ -0,0 +1,33 @@
+TASK_DESCRIPTION = """
+This 2-class linear classifier training loop has bugs preventing convergence.
+Fix it so that after 50 steps the loss is below 0.75 and decreasing.
+Model: nn.Linear(10, 2), dataset: random 2-class, 32 samples/batch.
+Print losses as: LOSSES:[val1, val2, ...]
+"""
+
+BUGGY_CODE = """
+import torch
+import torch.nn as nn
+torch.manual_seed(0)
+model = nn.Linear(10, 2)
+optimizer = torch.optim.Adam(model.parameters(), lr=10.0)  # BUG 1: lr too high
+criterion = nn.CrossEntropyLoss()
+losses = []
+for step in range(50):
+    x = torch.randn(32, 10)
+    y = torch.randint(0, 2, (32,))
+    optimizer.zero_grad()
+    logits = model(x)
+    loss = criterion(logits, y)
+    optimizer.step()   # BUG 2: step before backward
+    loss.backward()    # BUG 3: backward after step
+    losses.append(loss.item())
+print('##METRICS_START##')
+print('LOSSES:' + str(losses))
+print('##METRICS_END##')
+"""
+
+GROUND_TRUTH_BUGS = [
+    "optimizer.step() called before loss.backward()",
+    "learning rate 10.0 should be ~0.001",
+]

server/tasks/task2_nan_loss.py

@@ -0,0 +1,32 @@
+TASK_DESCRIPTION = """
+This binary regression trainer produces NaN loss around step 15.
+Fix the numerical instability so loss stays finite for all 60 steps
+and the final loss is below 0.5.
+Print losses as: LOSSES:[val1, val2, ...]
+"""
+
+BUGGY_CODE = """
+import torch
+import torch.nn as nn
+torch.manual_seed(42)
+model = nn.Linear(16, 1)
+optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+losses = []
+for step in range(60):
+    x = torch.randn(64, 16)
+    y = torch.rand(64, 1)
+    optimizer.zero_grad()
+    pred = torch.sigmoid(model(x))
+    # BUG: log(pred) can be -inf when pred rounds to 0.0
+    loss = -torch.mean(y * torch.log(pred) + (1 - y) * torch.log(1 - pred))
+    loss.backward()
+    optimizer.step()
+    losses.append(loss.item())
+print('##METRICS_START##')
+print('LOSSES:' + str(losses))
+print('##METRICS_END##')
+"""
+
+GROUND_TRUTH_BUGS = [
+    "torch.log(pred) when pred can be 0.0 after sigmoid — use F.binary_cross_entropy or clamp",
+]

server/tasks/task3_oom_leakage.py

@@ -0,0 +1,50 @@
+TASK_DESCRIPTION = """
+This trainer has TWO independent bugs:
+1. A memory leak causing OOM crash before epoch 5 on CPU.
+2. Data leakage inflating validation accuracy.
+Fix both. After 20 epochs: val_acc > 0.70, no OOM, no suspicious early accuracy spike.
+Print as: VAL_ACCS:[v1,v2,...] and FINAL_LOSS:X.XX
+"""
+
+BUGGY_CODE = """
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset, random_split
+
+torch.manual_seed(42)
+X = torch.randn(1000, 20)
+y = (X[:, 0] > 0).float()
+# BUG 1: augmentation before split — val set gets augmented
+X = X + torch.randn_like(X) * 0.1
+train_ds, val_ds = random_split(TensorDataset(X, y), [800, 200])
+model = nn.Sequential(nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 1))
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
+criterion = nn.BCEWithLogitsLoss()
+train_losses, val_accs = [], []
+total_loss = torch.tensor(0.0)  # BUG 2: keeps computation graph alive
+for epoch in range(20):
+    model.train()
+    for xb, yb in DataLoader(train_ds, batch_size=32):
+        optimizer.zero_grad()
+        out = model(xb).squeeze()
+        loss = criterion(out, yb)
+        loss.backward()
+        optimizer.step()
+        total_loss = total_loss + loss  # BUG 2: graph accumulates
+    model.eval()
+    with torch.no_grad():
+        idx = val_ds.indices
+        xv, yv = X[idx], y[idx]
+        preds = (torch.sigmoid(model(xv)) > 0.5).float()
+        acc = (preds == yv).float().mean().item()
+    val_accs.append(round(acc, 4))
+print('##METRICS_START##')
+print('VAL_ACCS:' + str(val_accs))
+print('FINAL_LOSS:' + str(total_loss.item()))
+print('##METRICS_END##')
+"""
+
+GROUND_TRUTH_BUGS = [
+    "Augmentation applied before split — move after split, apply to train only",
+    "total_loss += loss retains graph — use total_loss += loss.item()",
+]

server/tasks/task4_wrong_loss.py

@@ -0,0 +1,62 @@
+TASK_DESCRIPTION = """
+This is a multi-label classification problem where each sample can have multiple active classes.
+However, the model is currently using `CrossEntropyLoss`, which is meant for single-label (mutually exclusive) classes.
+Because of this, the loss trains but the predictions are essentially garbage (treating it as a single-label problem).
+
+Fix the loss function so it correctly handles multi-label classification.
+The grader will check:
+1. Loss convergence (loss < 0.5)
+2. Model predictions are actually multi-hot (avg > 1 label/sample)
+3. F1 Score > 0.6
+"""
+
+BUGGY_CODE = """
+import torch
+import torch.nn as nn
+from sklearn.metrics import f1_score
+
+torch.manual_seed(42)
+
+# Generate synthetic multi-label data (100 samples, 20 features, 5 classes)
+X = torch.randn(100, 20)
+# Each sample has a 30% chance of having each class active
+y = (torch.rand(100, 5) > 0.7).float() 
+
+model = nn.Sequential(
+    nn.Linear(20, 64),
+    nn.ReLU(),
+    nn.Linear(64, 5)
+)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+
+# BUG: CrossEntropyLoss is for single-label classification
+criterion = nn.CrossEntropyLoss()
+
+losses = []
+for step in range(100):
+    optimizer.zero_grad()
+    logits = model(X)
+    
+    # CrossEntropyLoss expects class indices, not one-hot/multi-hot vectors for the target
+    loss = criterion(logits, y)
+    
+    loss.backward()
+    optimizer.step()
+    losses.append(loss.item())
+
+# Evaluation
+with torch.no_grad():
+    logits = model(X)
+    # Using sigmoid and 0.5 threshold for multi-label prediction
+    preds = (torch.sigmoid(logits) > 0.5).float()
+    
+    avg_labels = preds.sum(dim=1).mean().item()
+    f1 = f1_score(y.numpy(), preds.numpy(), average='micro')
+
+print('##METRICS_START##')
+print('FINAL_LOSS:' + str(losses[-1]))
+print('AVG_LABELS:' + str(avg_labels))
+print('F1_SCORE:' + str(f1))
+print('##METRICS_END##')
+"""

server/tasks/task5_frozen_backbone.py

@@ -0,0 +1,62 @@
+TASK_DESCRIPTION = """
+This is a standard transfer learning setup classifying 10 categories.
+The developer froze the backbone during testing, but forgot to unfreeze it while still passing its parameters to the optimizer.
+Fix the code so the backbone actually trains, or only pass the head parameters.
+The grader checks the gradient norm of the backbone from the first backward pass.
+"""
+
+BUGGY_CODE = """
+import torch
+import torch.nn as nn
+
+torch.manual_seed(42)
+
+# Dummy dataset
+X = torch.randn(32, 512)
+y = torch.randint(0, 10, (32,))
+
+# A simulated pre-trained backbone
+backbone = nn.Sequential(
+    nn.Linear(512, 512),
+    nn.ReLU(),
+    nn.Linear(512, 512),
+    nn.ReLU()
+)
+
+# BUG: backbone is frozen, but passed to optimizer
+backbone.requires_grad_(False)
+
+head = nn.Linear(512, 10)
+
+# passing both backbone and head to optimizer even though backbone is frozen
+optimizer = torch.optim.Adam(
+    list(backbone.parameters()) + list(head.parameters()), lr=0.001
+)
+criterion = nn.CrossEntropyLoss()
+
+losses = []
+
+# Take one step to check gradients
+optimizer.zero_grad()
+features = backbone(X)
+logits = head(features)
+
+loss = criterion(logits, y)
+loss.backward()
+
+# Calculate gradient norm on backbone to see if it's training
+backbone_grad_norm = sum(
+    p.grad.norm().item() for p in backbone.parameters() if p.grad is not None
+)
+
+optimizer.step()
+losses.append(loss.item())
+
+# Note: if backbone is properly frozen and only head is passed, backbone_grad_norm will be 0 but optimizer won't complain.
+# If backbone is unfrozen, backbone_grad_norm will be > 0.
+# The grader handles both valid solutions.
+print('##METRICS_START##')
+print('FINAL_LOSS:' + str(losses[-1]))
+print('BACKBONE_GRAD_NORM:' + str(backbone_grad_norm))
+print('##METRICS_END##')
+"""