import requests import json import time BASE = "https://rak2315-ml-debug-env.hf.space" session = requests.Session() EPISODES = [ { "task_id": "shape_mismatch", "label": "Episode 1: Shape Mismatch (Easy)", "steps": [ {"action_type": "inspect", "tool_name": "run_code"}, {"action_type": "inspect", "tool_name": "get_traceback"}, { "action_type": "fix", "bug_type": "shape_mismatch", "diagnosis": "nn.Linear input dimension wrong — fc2 expects 128 but fc1 outputs 64", "fixed_code": """ import torch import torch.nn as nn class Model(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(64, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = torch.relu(self.fc1(x)) return self.fc2(x) model = Model() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) loss_fn = nn.CrossEntropyLoss() for epoch in range(3): x = torch.randn(32, 64) y = torch.randint(0, 10, (32,)) optimizer.zero_grad() out = model(x) loss = loss_fn(out, y) loss.backward() optimizer.step() print(f"Epoch {epoch+1}, loss: {loss.item():.4f}") print("Training finished") """.strip(), }, ], }, { "task_id": "gradient_not_zeroed", "label": "Episode 2: Gradient Not Zeroed (Medium-Hard)", "steps": [ {"action_type": "inspect", "tool_name": "inspect_gradients"}, { "action_type": "fix", "bug_type": "gradient_not_zeroed", "diagnosis": "optimizer.zero_grad() missing before loss.backward() — gradients accumulate across batches causing explosion", "fixed_code": """ import torch import torch.nn as nn model = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 1)) optimizer = torch.optim.SGD(model.parameters(), lr=0.01) loss_fn = nn.MSELoss() for epoch in range(5): x = torch.randn(32, 16) y = torch.randn(32, 1) optimizer.zero_grad() out = model(x) loss = loss_fn(out, y) loss.backward() optimizer.step() print(f"Epoch {epoch+1}, loss: {loss.item():.4f}") print("Training finished") """.strip(), }, ], }, { "task_id": "compound_leakage_eval", "label": "Episode 3: Compound Leakage + Eval Mode (Expert — 2 bugs)", "steps": [ {"action_type": "inspect", "tool_name": "run_code"}, {"action_type": "inspect", "tool_name": "print_shapes"}, { "action_type": "fix", "bug_type": "compound_leakage_eval", "diagnosis": "Two bugs: (1) normalization computed on full dataset before train/test split causes data leakage, (2) model.eval() missing during evaluation causes non-deterministic metrics due to active dropout", "fixed_code": """ import torch import torch.nn as nn torch.manual_seed(42) X = torch.randn(200, 16) y = (X[:, 0] > 0).float() split = 160 X_train, X_test = X[:split], X[split:] y_train, y_test = y[:split], y[split:] mean = X_train.mean(dim=0) std = X_train.std(dim=0) + 1e-8 X_train = (X_train - mean) / std X_test = (X_test - mean) / std model = nn.Sequential(nn.Linear(16, 32), nn.Dropout(0.3), nn.ReLU(), nn.Linear(32, 1)) optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) loss_fn = nn.BCEWithLogitsLoss() model.train() for epoch in range(5): optimizer.zero_grad() out = model(X_train).squeeze() loss = loss_fn(out, y_train) loss.backward() optimizer.step() print(f"Epoch {epoch+1}, loss: {loss.item():.4f}") model.eval() with torch.no_grad(): preds = (model(X_test).squeeze() > 0).float() acc = (preds == y_test).float().mean().item() print(f"Accuracy: {acc:.4f}") print("Evaluation complete") """.strip(), }, ], }, ] def print_separator(char="-", width=60): print(char * width) def run_episode(episode): task_id = episode["task_id"] label = episode["label"] print_separator("=") print(f" {label}") print_separator("=") r = session.post(f"{BASE}/reset", json={"task_id": task_id}) if r.status_code != 200: print(f" Reset failed: {r.status_code}") return obs = r.json()["observation"] print(f"\n Alert: \"{obs['alert']}\"") print(f" Tools: {obs['available_tools']}") print(f" Step budget: {obs['step_budget']} | Bugs: {obs['num_bugs']}") print() for i, step in enumerate(episode["steps"]): action_type = step["action_type"] time.sleep(0.5) if action_type == "inspect": tool = step["tool_name"] r = session.post(f"{BASE}/step", json={"action": step}) if r.status_code != 200: print(f" Step {i+1}: inspect:{tool} → ERROR {r.status_code}") continue obs = r.json()["observation"] result = obs.get("tool_result", "") or "" preview = result[:120].replace("\n", " ").strip() budget = obs.get("step_budget", "?") print(f" Step {i+1}: inspect:{tool:<20} reward=+0.00 budget={budget}") print(f" → {preview}...") print() elif action_type == "fix": r = session.post(f"{BASE}/step", json={"action": step}) if r.status_code != 200: print(f" Step {i+1}: fix → ERROR {r.status_code}") continue obs = r.json()["observation"] score = obs.get("grader_score", 0) or 0 feedback = obs.get("grader_feedback", "") or "" multiplier = obs.get("efficiency_multiplier", 1.0) or 1.0 budget = obs.get("step_budget", "?") status = "✅ FIXED" if score >= 0.95 else ("🟡 PARTIAL" if score >= 0.6 else "❌") print(f" Step {i+1}: fix:{step['bug_type']:<25} reward={score:.2f} {status}") if multiplier > 1.0: print(f" → Efficiency bonus: ×{multiplier} applied") print(f" → {feedback[:120]}") print() print() def main(): print() print("=" * 60) print(" ML Debug Env — Live Demo") print(" Agent debugs broken PyTorch scripts using tool calls") print(" Partial observability: alert only on reset, no code") print("=" * 60) print() print(f" Environment: {BASE}") r = session.get(f"{BASE}/health") print(f" Health: {r.json()}") print() for episode in EPISODES: run_episode(episode) time.sleep(1) print_separator("=") print(" Demo complete.") print(" Scoring ladder: 0.01 → wrong type | 0.20 → crashes") print(" 0.40 → incomplete | 0.60 → not fixed") print(" 0.80 → missing signal | 0.99 → perfect ✅") print_separator("=") print() if __name__ == "__main__": main()