feat: complete files

.gitignore

@@ -0,0 +1,7 @@
+__pycache__/
+*.pyc
+*.pyo
+.env
+.pytest_cache/
+dist/
+*.egg-info/

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.11-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PORT=7860
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+EXPOSE 7860
+CMD ["uvicorn", "src.pytorch_debug_env.server:app", "--host", "0.0.0.0", "--port", "7860"]

inference.py

@@ -0,0 +1,126 @@
+# inference.py
+import asyncio
+import json
+import os
+from typing import List
+
+from openai import OpenAI
+import httpx
+
+API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
+MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-3.5-turbo")
+API_KEY = os.environ.get("OPENAI_API_KEY", "dummy")
+ENV_URL = os.environ.get("ENV_URL", "http://localhost:7860")
+TASK_NAME = os.environ.get("TASK_NAME", "easy")
+MAX_STEPS = int(os.environ.get("MAX_STEPS", "5"))
+SUCCESS_SCORE_THRESHOLD = float(os.environ.get("SUCCESS_SCORE_THRESHOLD", "0.7"))
+MAX_TOTAL_REWARD = float(os.environ.get("MAX_TOTAL_REWARD", "1.0"))
+
+
+def log_start(task, env, model):
+    print(json.dumps({
+        "type": "START",
+        "task": task,
+        "env": env,
+        "model": model,
+    }), flush=True)
+
+
+def log_step(step, action, reward, done, error):
+    print(json.dumps({
+        "type": "STEP",
+        "step": step,
+        "action": action,
+        "reward": float(reward),
+        "done": bool(done),
+        "error": error,
+    }), flush=True)
+
+
+def log_end(success, steps, score, rewards):
+    print(json.dumps({
+        "type": "END",
+        "success": bool(success),
+        "steps": steps,
+        "score": float(score),
+        "rewards": [float(r) for r in rewards],
+    }), flush=True)
+
+
+def get_model_message(client: OpenAI, observation: dict, history: List[str]) -> str:
+    prompt = f"""
+You are debugging a PyTorch training job. Respond ONLY with valid JSON matching this exact schema:
+{{
+  "current_hypothesis": {{"bug_type": "<string>", "affected_file": "<string>", "confidence": <0.0-1.0>}},
+  "investigation_action": {{"action": "reveal_file", "target": "<filename>"}},
+  "commit_diagnosis": false,
+  "final_diagnosis": null
+}}
+
+Valid action types: reveal_file, extend_loss_curve, extend_gpu_profile, reveal_log_chunk, run_diagnostic
+Valid bug types: missing_zero_grad, data_leakage, memory_leak, learning_rate_too_high, gradient_explosion
+
+Observation:
+{json.dumps(observation)[:8000]}
+History: {history}
+"""
+    completion = client.chat.completions.create(
+        model=MODEL_NAME,
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0,
+        max_tokens=500,
+    )
+    return (completion.choices[0].message.content or "").strip()
+
+
+async def main():
+    client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY)
+    rewards = []
+    history = []
+    steps_taken = 0
+    score = 0.0
+    success = False
+
+    log_start(task=TASK_NAME, env="pytorch-debug-env", model=MODEL_NAME)
+
+    async with httpx.AsyncClient(timeout=60.0) as session:
+        reset_resp = await session.post(f"{ENV_URL}/reset", params={"task_id": TASK_NAME})
+        reset_resp.raise_for_status()
+        result = reset_resp.json()
+        session_id = result.get("session_id")
+        observation = result["observation"]
+
+        for step in range(1, MAX_STEPS + 1):
+            if result.get("done"):
+                break
+
+            action_text = get_model_message(client, observation, history)
+            try:
+                action_json = json.loads(action_text)
+                step_resp = await session.post(f"{ENV_URL}/step", params={"session_id": session_id}, json=action_json)
+                step_resp.raise_for_status()
+                result = step_resp.json()
+                reward = result.get("reward", 0.0)
+                done = result.get("done", False)
+                error = None
+                observation = result["observation"]
+            except Exception as exc:
+                reward = 0.0
+                done = True
+                error = str(exc)
+
+            rewards.append(reward)
+            steps_taken = step
+            log_step(step=step, action=action_text, reward=reward, done=done, error=error)
+            history.append(f"step={step} reward={reward:.3f}")
+
+            if done:
+                break
+
+    score = min(max(rewards[-1] if rewards else 0.0, 0.0), 1.0)
+    success = score >= SUCCESS_SCORE_THRESHOLD
+    log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

openenv.yaml

@@ -0,0 +1,40 @@
+name: pytorch-debug-env
+version: 1.0.0
+description: Multi-step OpenEnv environment for diagnosing broken PyTorch training jobs.
+author: Priyansh Saxena
+
+client:
+  class_name: PyTorchDebugEnv
+  module: src.pytorch_debug_env.environment
+
+action:
+  class_name: PyTorchDebugAction
+  module: src.pytorch_debug_env.models
+
+observation:
+  class_name: PyTorchDebugObservation
+  module: src.pytorch_debug_env.models
+
+default_image: pytorch-debug-env:latest
+spec_version: 1
+
+tags:
+  - openenv
+  - pytorch
+  - debugging
+  - reinforcement-learning
+
+tasks:
+  - id: easy
+    name: Single-file bug detection
+    difficulty: easy
+  - id: medium
+    name: Multi-file root cause analysis
+    difficulty: medium
+  - id: hard
+    name: Silent failure diagnosis
+    difficulty: hard
+
+runtime:
+  framework: fastapi
+  container_port: 7860

pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+asyncio_mode = auto

requirements.txt

@@ -0,0 +1,9 @@
+fastapi==0.115.0
+uvicorn[standard]==0.30.6
+pydantic==2.9.2
+numpy==2.1.1
+openai==1.51.0
+httpx==0.27.2
+pytest==8.3.3
+pytest-asyncio==0.24.0
+openenv>=0.1.0

scenarios/seeds.json

@@ -0,0 +1 @@
+{}

src/pytorch_debug_env/bug_library.py

@@ -0,0 +1,140 @@
+# src/pytorch_debug_env/bug_library.py
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional
+import numpy as np
+
+
+@dataclass
+class BugTemplate:
+    bug_type: str
+    category: str
+    difficulty: str
+    primary_bug_file: str
+    related_files: List[str]
+    red_herring_file: Optional[str]
+    fix_strategy: str
+    line_range: List[int]
+    description: str
+    artifact_generator: Callable
+    repo_mutator: Callable
+    metadata: Dict = field(default_factory=dict)
+
+
+BUG_CATEGORIES = {
+    "shape_mismatch": "model",
+    "missing_zero_grad": "optimization",
+    "wrong_loss_function": "optimization",
+    "learning_rate_too_high": "optimization",
+    "gradient_explosion": "optimization",
+    "memory_leak": "resource",
+    "data_leakage": "data",
+    "incorrect_normalization": "data",
+    "distributed_sync_error": "distributed",
+    "amp_overflow": "numerics",
+}
+
+# Realistic artifact generator
+def dummy_artifact_generator(artifact_type: str, rng):
+    if artifact_type == "loss_curve":
+        t = np.arange(100)
+        base = 2.3 * np.exp(-0.01 * t) + 0.15
+        oscillation = 0.22 * np.sin(0.25 * t) * np.exp(-0.002 * t)
+        return [
+            {"step": int(i), "train_loss": float(base[i] + oscillation[i])}
+            for i in range(100)
+        ]
+    elif artifact_type == "gpu_profile":
+        t = np.arange(100)
+        allocated = 2048 + 2.4 * t
+        return [
+            {"step": int(i), "allocated_mb": float(allocated[i])}
+            for i in range(100)
+        ]
+    elif artifact_type == "training_log":
+        return "Epoch 1, Step 0: loss 2.45\nEpoch 1, Step 1: loss 2.43\n"
+    return []
+
+def mutate_missing_zero_grad(repo_files, rng):
+    repo_files["train.py"] = """import torch
+from model.architecture import Net
+
+model = Net()
+optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+criterion = torch.nn.CrossEntropyLoss()
+
+for epoch in range(10):
+    for x, y in dataloader:
+        # optimizer.zero_grad()  # BUG: commented out
+        output = model(x)
+        loss = criterion(output, y)
+        loss.backward()
+        optimizer.step()
+"""
+    return repo_files
+
+def mutate_data_leakage(repo_files, rng):
+    repo_files["data/dataset.py"] = """from torch.utils.data import Dataset
+
+class ImageDataset(Dataset):
+    def __init__(self, data, split="train"):
+        # BUG: We use the entire data instead of just the split
+        self.data = data
+        self.split = split
+"""
+    return repo_files
+
+def mutate_memory_leak(repo_files, rng):
+    repo_files["data/dataset.py"] = """from torch.utils.data import Dataset
+
+class ImageDataset(Dataset):
+    def __init__(self):
+        # BUG: Storing huge tensors in a class-level variable leading to memory accumulation
+        self.cache = []
+
+    def load(self, x):
+        self.cache.append(x)
+        return x
+"""
+    return repo_files
+
+BUG_TEMPLATES = [
+    BugTemplate(
+        bug_type="missing_zero_grad",
+        category="optimization",
+        difficulty="easy",
+        primary_bug_file="train.py",
+        related_files=[],
+        red_herring_file="model/architecture.py",
+        fix_strategy="Call optimizer.zero_grad() before loss.backward()",
+        line_range=[9, 14],
+        description="Missing zero grad",
+        artifact_generator=dummy_artifact_generator,
+        repo_mutator=mutate_missing_zero_grad,
+    ),
+    BugTemplate(
+        bug_type="data_leakage",
+        category="data",
+        difficulty="medium",
+        primary_bug_file="data/dataset.py",
+        related_files=["data/preprocessing.py"],
+        red_herring_file="train.py",
+        fix_strategy="Ensure validation split is strictly separate from training",
+        line_range=[4, 6],
+        description="Data leakage",
+        artifact_generator=dummy_artifact_generator,
+        repo_mutator=mutate_data_leakage,
+    ),
+    BugTemplate(
+        bug_type="memory_leak",
+        category="resource",
+        difficulty="hard",
+        primary_bug_file="data/dataset.py",
+        related_files=["train.py"],
+        red_herring_file="model/attention.py",
+        fix_strategy="Avoid holding reference to tensors in class cache",
+        line_range=[5, 9],
+        description="Memory leak",
+        artifact_generator=dummy_artifact_generator,
+        repo_mutator=mutate_memory_leak,
+    )
+]

src/pytorch_debug_env/environment.py

@@ -0,0 +1,168 @@
+# src/pytorch_debug_env/environment.py
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+from .models import (
+    HypothesisRecord,
+    PyTorchDebugAction,
+    PyTorchDebugObservation,
+    PyTorchDebugState,
+)
+from .reward import compute_step_reward
+from .scenario_generator import ScenarioGenerator
+from .graders import grade_easy, grade_medium, grade_hard
+
+GRADER_MAP = {"easy": grade_easy, "medium": grade_medium, "hard": grade_hard}
+
+
+@dataclass
+class RuntimeState:
+    scenario: object | None = None
+    max_steps: int = 5
+    current_step: int = 0
+    revealed_files: List[str] = field(default_factory=list)
+    hypothesis_history: List[HypothesisRecord] = field(default_factory=list)
+    done: bool = False
+    final_score: float = 0.0
+
+
+class PyTorchDebugEnv:
+    def __init__(self, generator: ScenarioGenerator, max_steps: int = 5):
+        self.generator = generator
+        self.runtime = RuntimeState(max_steps=max_steps)
+
+    async def reset(self, task_id: str = "easy"):
+        scenario = self.generator.generate(task_id)
+        self.runtime = RuntimeState(
+            scenario=scenario,
+            max_steps=5 if task_id == "easy" else 6,
+            current_step=0,
+            revealed_files=["train.py", "config/training_config.yaml"],
+            hypothesis_history=[],
+            done=False,
+            final_score=0.0,
+        )
+        return self._build_observation(last_feedback="Episode reset.")
+
+    async def step(self, action: PyTorchDebugAction):
+        if self.runtime.scenario is None:
+            raise RuntimeError("Call /reset before /step")
+
+        if self.runtime.done:
+            raise RuntimeError("Episode already completed")
+
+        self.runtime.current_step += 1
+        scenario = self.runtime.scenario
+        previous_quality = self.runtime.hypothesis_history[-1].quality if self.runtime.hypothesis_history else 0.0
+
+        investigation_target = None
+        if action.investigation_action and action.investigation_action.action == "reveal_file":
+            investigation_target = action.investigation_action.target
+            if investigation_target in scenario.repo_files and investigation_target not in self.runtime.revealed_files:
+                self.runtime.revealed_files.append(investigation_target)
+
+        committed = action.final_diagnosis.model_dump() if action.commit_diagnosis and action.final_diagnosis else None
+        reward, components = compute_step_reward(
+            previous_quality=previous_quality,
+            current_hypothesis=action.current_hypothesis.model_dump(),
+            ground_truth=scenario.ground_truth,
+            investigation_target=investigation_target,
+            committed_diagnosis=None, # Temporarily don't compute diagnosis reward here to use grader
+            step_num=self.runtime.current_step,
+            max_steps=self.runtime.max_steps,
+        )
+
+        if committed:
+            grader = GRADER_MAP.get(scenario.task_id, grade_easy)
+            diagnosis_reward = grader(committed, scenario.ground_truth)
+
+            # Combine the diagnosis reward logic from `compute_step_reward` that applies on top
+            if diagnosis_reward > 0.7:
+                diagnosis_reward += max(0.0, 0.08 * (self.runtime.max_steps - self.runtime.current_step))
+
+            # Update the total reward incorporating diagnosis
+            components["diagnosis_reward"] = round(diagnosis_reward, 4)
+            delta = components["hypothesis_delta"]
+            inv_reward = components["investigation_reward"]
+            conf_bonus = components["confirmation_bonus"]
+
+            total = 0.60 * delta + 0.20 * inv_reward + 0.20 * diagnosis_reward + conf_bonus
+            reward = round(min(max(total, 0.0), 1.0), 4)
+
+        self.runtime.hypothesis_history.append(
+            HypothesisRecord(
+                step=self.runtime.current_step,
+                hypothesis=action.current_hypothesis,
+                quality=components["hypothesis_quality"],
+            )
+        )
+
+        if action.commit_diagnosis or self.runtime.current_step >= self.runtime.max_steps:
+            self.runtime.done = True
+            self.runtime.final_score = reward
+
+        observation = self._build_observation(
+            last_feedback=self._feedback(action, scenario.ground_truth)
+        )
+        return {
+            "observation": observation,
+            "reward": reward,
+            "done": self.runtime.done,
+            "info": components,
+        }
+
+    async def state(self):
+        scenario = self.runtime.scenario
+        if not scenario:
+            return None
+        return PyTorchDebugState(
+            scenario_id=scenario.scenario_id,
+            task_id=scenario.task_id,
+            max_steps=self.runtime.max_steps,
+            current_step=self.runtime.current_step,
+            revealed_files=self.runtime.revealed_files,
+            remaining_files=[
+                f for f in scenario.repo_files.keys() if f not in self.runtime.revealed_files
+            ],
+            done=self.runtime.done,
+            final_score=self.runtime.final_score,
+        )
+
+    def _build_observation(self, last_feedback: str) -> PyTorchDebugObservation:
+        scenario = self.runtime.scenario
+        revealed = {k: v for k, v in scenario.repo_files.items() if k in self.runtime.revealed_files}
+        available = [k for k in scenario.repo_files.keys() if k not in self.runtime.revealed_files]
+
+        loss_window_size = min(len(scenario.loss_curve), 100 * (self.runtime.current_step + 1))
+        gpu_window_size = min(len(scenario.gpu_profile), 100 * (self.runtime.current_step + 1))
+        log_lines = scenario.training_log.splitlines()
+        visible_log = "\n".join(log_lines[-min(len(log_lines), 10 * (self.runtime.current_step + 1)):])
+
+        return PyTorchDebugObservation(
+            scenario_id=scenario.scenario_id,
+            task_id=scenario.task_id,
+            revealed_files=revealed,
+            available_files=available,
+            loss_curve_window=scenario.loss_curve[:loss_window_size],
+            gpu_profile_window=scenario.gpu_profile[:gpu_window_size],
+            training_log_tail=visible_log,
+            step_num=self.runtime.current_step,
+            steps_remaining=max(0, self.runtime.max_steps - self.runtime.current_step),
+            investigation_budget=max(0, self.runtime.max_steps - self.runtime.current_step),
+            hypothesis_history=self.runtime.hypothesis_history,
+            last_feedback=last_feedback,
+        )
+
+    def _feedback(self, action: PyTorchDebugAction, gt: Dict) -> str:
+        suspected_file = action.current_hypothesis.affected_file
+        suspected_bug = action.current_hypothesis.bug_type
+
+        if suspected_file == gt.get("red_herring_file"):
+            return "That file contains a plausible symptom, but not the root cause. Investigate upstream causes."
+        if suspected_file == gt["primary_bug_file"] and suspected_bug != gt["bug_type"]:
+            return "Correct region, wrong failure mode. Re-check the training artifacts more carefully."
+        if suspected_bug == gt["bug_type"] and suspected_file != gt["primary_bug_file"]:
+            return "The bug class looks right, but the faulty implementation is in another file."
+        return "Continue refining the hypothesis using newly revealed evidence."

src/pytorch_debug_env/graders.py

@@ -0,0 +1,31 @@
+# src/pytorch_debug_env/graders.py
+from __future__ import annotations
+
+from .reward import final_diagnosis_score
+
+
+def grade_easy(action: dict, gt: dict) -> float:
+    return final_diagnosis_score(action, gt)
+
+
+def grade_medium(action: dict, gt: dict) -> float:
+    score = final_diagnosis_score(action, gt)
+    if action.get("affected_file") in gt.get("related_files", []):
+        score = min(1.0, score + 0.05)
+    return round(score, 4)
+
+
+def grade_hard(action: dict, gt: dict) -> float:
+    score = final_diagnosis_score(action, gt)
+
+    # partial credit if model gets the right category on subtle bugs
+    if score < 0.2 and action.get("bug_type"):
+        if gt.get("category"):
+            from .bug_library import BUG_CATEGORIES
+            if BUG_CATEGORIES.get(action["bug_type"]) == gt["category"]:
+                score = max(score, 0.18)
+
+    if action.get("affected_file") == gt.get("red_herring_file"):
+        score = max(0.0, score - 0.1)
+
+    return round(min(score, 1.0), 4)

src/pytorch_debug_env/models.py

@@ -0,0 +1,74 @@
+# src/pytorch_debug_env/models.py
+from __future__ import annotations
+
+from typing import Dict, List, Literal, Optional
+from pydantic import BaseModel, Field
+
+
+class Hypothesis(BaseModel):
+    bug_type: str = Field(..., description="Current suspected bug type")
+    affected_file: str = Field(..., description="Current suspected file")
+    confidence: float = Field(..., ge=0.0, le=1.0)
+
+
+class InvestigationAction(BaseModel):
+    action: Literal[
+        "reveal_file",
+        "extend_loss_curve",
+        "extend_gpu_profile",
+        "reveal_log_chunk",
+        "run_diagnostic",
+    ]
+    target: Optional[str] = None
+
+
+class FinalDiagnosis(BaseModel):
+    bug_type: str
+    affected_file: str
+    line_range: List[int]
+    fix_strategy: str
+    confidence: float = Field(..., ge=0.0, le=1.0)
+
+
+class PyTorchDebugAction(BaseModel):
+    current_hypothesis: Hypothesis
+    investigation_action: Optional[InvestigationAction] = None
+    commit_diagnosis: bool = False
+    final_diagnosis: Optional[FinalDiagnosis] = None
+
+
+class HypothesisRecord(BaseModel):
+    step: int
+    hypothesis: Hypothesis
+    quality: float
+
+
+class PyTorchDebugObservation(BaseModel):
+    scenario_id: str
+    task_id: str
+    revealed_files: Dict[str, str]
+    available_files: List[str]
+    loss_curve_window: List[Dict]
+    gpu_profile_window: List[Dict]
+    training_log_tail: str
+    step_num: int
+    steps_remaining: int
+    investigation_budget: int
+    hypothesis_history: List[HypothesisRecord]
+    last_feedback: str
+
+
+class PyTorchDebugState(BaseModel):
+    scenario_id: str
+    task_id: str
+    max_steps: int
+    current_step: int
+    revealed_files: List[str]
+    remaining_files: List[str]
+    done: bool
+    final_score: float = 0.0
+
+
+class PyTorchDebugReward(BaseModel):
+    value: float = Field(..., ge=0.0, le=1.0)
+    components: Dict[str, float]

src/pytorch_debug_env/reward.py

@@ -0,0 +1,92 @@
+# src/pytorch_debug_env/reward.py
+from __future__ import annotations
+
+from .bug_library import BUG_CATEGORIES
+
+
+def hypothesis_quality(hypothesis: dict, ground_truth: dict) -> float:
+    q = 0.0
+
+    if hypothesis.get("affected_file") == ground_truth["primary_bug_file"]:
+        q += 0.45
+    elif hypothesis.get("affected_file") in ground_truth.get("related_files", []):
+        q += 0.15
+
+    if hypothesis.get("bug_type") == ground_truth["bug_type"]:
+        q += 0.40
+    elif BUG_CATEGORIES.get(hypothesis.get("bug_type")) == BUG_CATEGORIES.get(ground_truth["bug_type"]):
+        q += 0.13
+
+    calibration = 1.0 - abs(hypothesis.get("confidence", 0.5) - min(q, 1.0))
+    q += 0.15 * calibration
+    return round(min(q, 1.0), 4)
+
+
+def final_diagnosis_score(diagnosis: dict, ground_truth: dict) -> float:
+    score = 0.0
+
+    if diagnosis.get("bug_type") == ground_truth["bug_type"]:
+        score += 0.40
+    if diagnosis.get("affected_file") == ground_truth["primary_bug_file"]:
+        score += 0.25
+
+    predicted = diagnosis.get("line_range", [0, 0])
+    actual = ground_truth.get("line_range", [0, 0])
+    overlap = line_overlap(predicted, actual)
+    score += 0.20 * overlap
+
+    if diagnosis.get("fix_strategy") == ground_truth["fix_strategy"]:
+        score += 0.15
+
+    return round(min(score, 1.0), 4)
+
+
+def line_overlap(pred: list[int], actual: list[int]) -> float:
+    p1, p2 = pred
+    a1, a2 = actual
+    inter = max(0, min(p2, a2) - max(p1, a1) + 1)
+    union = max(p2, a2) - min(p1, a1) + 1
+    return inter / union if union else 0.0
+
+
+def compute_step_reward(
+    previous_quality: float,
+    current_hypothesis: dict,
+    ground_truth: dict,
+    investigation_target: str | None = None,
+    committed_diagnosis: dict | None = None,
+    step_num: int = 1,
+    max_steps: int = 5,
+) -> tuple[float, dict]:
+    current_quality = hypothesis_quality(current_hypothesis, ground_truth)
+    delta = current_quality - previous_quality
+
+    confirmation_bonus = 0.03 * current_quality if abs(delta) < 0.01 else 0.0
+
+    investigation_reward = 0.0
+    if investigation_target:
+        if investigation_target == ground_truth["primary_bug_file"]:
+            investigation_reward = 0.07
+        elif investigation_target in ground_truth.get("related_files", []):
+            investigation_reward = 0.025
+        elif investigation_target == ground_truth.get("red_herring_file"):
+            investigation_reward = -0.04
+        else:
+            investigation_reward = -0.01
+
+    diagnosis_reward = 0.0
+    if committed_diagnosis:
+        diagnosis_reward = final_diagnosis_score(committed_diagnosis, ground_truth)
+        if diagnosis_reward > 0.7:
+            diagnosis_reward += max(0.0, 0.08 * (max_steps - step_num))
+
+    total = 0.60 * delta + 0.20 * investigation_reward + 0.20 * diagnosis_reward + confirmation_bonus
+    total = round(min(max(total, -0.2), 1.0), 4)
+
+    return total, {
+        "hypothesis_quality": current_quality,
+        "hypothesis_delta": round(delta, 4),
+        "investigation_reward": round(investigation_reward, 4),
+        "diagnosis_reward": round(diagnosis_reward, 4),
+        "confirmation_bonus": round(confirmation_bonus, 4),
+    }

src/pytorch_debug_env/scenario_generator.py

@@ -0,0 +1,86 @@
+# src/pytorch_debug_env/scenario_generator.py
+from __future__ import annotations
+
+import random
+import uuid
+from dataclasses import dataclass
+from typing import Dict, List
+
+import numpy as np
+
+from .bug_library import BugTemplate
+
+
+@dataclass
+class Scenario:
+    scenario_id: str
+    task_id: str
+    repo_files: Dict[str, str]
+    loss_curve: List[Dict]
+    gpu_profile: List[Dict]
+    training_log: str
+    ground_truth: Dict
+
+
+class ScenarioGenerator:
+    def __init__(self, bug_templates: List[BugTemplate]):
+        self.bug_templates = bug_templates
+
+    def generate(self, difficulty: str, seed: int | None = None) -> Scenario:
+        rng = random.Random(seed)
+        template = rng.choice([b for b in self.bug_templates if b.difficulty == difficulty])
+
+        repo_files = self._base_repo(rng)
+        repo_files = template.repo_mutator(repo_files, rng)
+
+        loss_curve = template.artifact_generator("loss_curve", rng)
+        gpu_profile = template.artifact_generator("gpu_profile", rng)
+        training_log = template.artifact_generator("training_log", rng)
+
+        ground_truth = {
+            "bug_type": template.bug_type,
+            "category": template.category,
+            "primary_bug_file": template.primary_bug_file,
+            "related_files": template.related_files,
+            "red_herring_file": template.red_herring_file,
+            "fix_strategy": template.fix_strategy,
+            "line_range": template.line_range,
+        }
+
+        return Scenario(
+            scenario_id=str(uuid.uuid4())[:8],
+            task_id=difficulty,
+            repo_files=repo_files,
+            loss_curve=loss_curve,
+            gpu_profile=gpu_profile,
+            training_log=training_log,
+            ground_truth=ground_truth,
+        )
+
+    def _base_repo(self, rng: random.Random) -> Dict[str, str]:
+        return {
+            "train.py": self._train_py(),
+            "model/architecture.py": self._model_py(),
+            "model/attention.py": self._attention_py(),
+            "data/dataset.py": self._dataset_py(),
+            "data/preprocessing.py": self._preprocess_py(),
+            "config/training_config.yaml": self._config_yaml(),
+        }
+
+    def _train_py(self) -> str:
+        return """import torch\nfrom model.architecture import Net\n\n# training loop placeholder\n"""
+
+    def _model_py(self) -> str:
+        return """import torch.nn as nn\n\nclass Net(nn.Module):\n    def __init__(self):\n        super().__init__()\n"""
+
+    def _attention_py(self) -> str:
+        return """# custom attention layer\n"""
+
+    def _dataset_py(self) -> str:
+        return """from torch.utils.data import Dataset\n\nclass ImageDataset(Dataset):\n    pass\n"""
+
+    def _preprocess_py(self) -> str:
+        return """def normalize(x):\n    return x\n"""
+
+    def _config_yaml(self) -> str:
+        return "lr: 0.001\nbatch_size: 32\n"

src/pytorch_debug_env/server.py

@@ -0,0 +1,55 @@
+# src/pytorch_debug_env/server.py
+from fastapi import FastAPI, Query
+from uuid import uuid4
+
+from .environment import PyTorchDebugEnv
+from .models import PyTorchDebugAction
+from .scenario_generator import ScenarioGenerator
+from .bug_library import BUG_TEMPLATES
+
+app = FastAPI(title="PyTorch Debug Env")
+
+sessions = {}
+latest_session_id = None
+
+@app.get("/")
+async def root():
+    return {
+        "name": "pytorch-debug-env",
+        "version": "1.0.0",
+        "endpoints": ["/reset", "/step", "/state", "/health"],
+        "tasks": ["easy", "medium", "hard"]
+    }
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+
+@app.post("/reset")
+async def reset(task_id: str = "easy"):
+    global latest_session_id
+    session_id = str(uuid4())
+    env = PyTorchDebugEnv(generator=ScenarioGenerator(BUG_TEMPLATES))
+    sessions[session_id] = env
+    latest_session_id = session_id
+    obs = await env.reset(task_id=task_id)
+    return {"session_id": session_id, "observation": obs, "done": False}
+
+
+@app.post("/step")
+async def step(action: PyTorchDebugAction, session_id: str = Query(None)):
+    sid = session_id or latest_session_id
+    env = sessions.get(sid)
+    if not env:
+        return {"error": "Invalid session_id"}
+    return await env.step(action)
+
+
+@app.get("/state")
+async def state(session_id: str = Query(None)):
+    sid = session_id or latest_session_id
+    env = sessions.get(sid)
+    if not env:
+        return {"error": "Invalid session_id"}
+    return await env.state()

tests/conftest.py

@@ -0,0 +1,5 @@
+import pytest
+
+@pytest.fixture
+def anyio_backend():
+    return "asyncio"

tests/test_environment.py

@@ -0,0 +1,12 @@
+# tests/test_environment.py
+import pytest
+from src.pytorch_debug_env.environment import PyTorchDebugEnv
+from src.pytorch_debug_env.scenario_generator import ScenarioGenerator
+from src.pytorch_debug_env.bug_library import BUG_TEMPLATES
+
+@pytest.mark.asyncio
+async def test_env_reset():
+    generator = ScenarioGenerator(BUG_TEMPLATES)
+    env = PyTorchDebugEnv(generator=generator)
+    obs = await env.reset("easy")
+    assert obs.task_id == "easy"

tests/test_graders.py

@@ -0,0 +1,19 @@
+# tests/test_graders.py
+from src.pytorch_debug_env.graders import grade_easy
+
+def test_grade_easy():
+    gt = {
+        "bug_type": "missing_zero_grad",
+        "primary_bug_file": "train.py",
+        "related_files": [],
+        "line_range": [10, 15],
+        "fix_strategy": "Call optimizer.zero_grad() before loss.backward()",
+    }
+    action = {
+        "bug_type": "missing_zero_grad",
+        "affected_file": "train.py",
+        "line_range": [10, 15],
+        "fix_strategy": "Call optimizer.zero_grad() before loss.backward()",
+        "confidence": 0.8
+    }
+    assert grade_easy(action, gt) > 0.8

tests/test_reward.py

@@ -0,0 +1,16 @@
+# tests/test_reward.py
+from src.pytorch_debug_env.reward import hypothesis_quality
+
+
+def test_hypothesis_quality_exact_match():
+    gt = {
+        "bug_type": "missing_zero_grad",
+        "primary_bug_file": "train.py",
+        "related_files": [],
+    }
+    hyp = {
+        "bug_type": "missing_zero_grad",
+        "affected_file": "train.py",
+        "confidence": 0.8,
+    }
+    assert hypothesis_quality(hyp, gt) > 0.8