Spaces:
Sleeping
Sleeping
File size: 5,274 Bytes
c14504c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """
Core environment logic: task loading, reset, step, state.
"""
import json
import random
import uuid
from pathlib import Path
from .executor import run_code_safely
class CodeDebugEnvironment:
def __init__(self):
self.tasks: dict[str, dict] = {}
self.episodes: dict[str, dict] = {}
self._load_tasks()
def _load_tasks(self):
tasks_dir = Path(__file__).parent.parent / "tasks"
for json_file in sorted(tasks_dir.rglob("*.json")):
with open(json_file, encoding="utf-8") as f:
task = json.load(f)
self.tasks[task["task_id"]] = task
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def reset(self, task_id: str | None = None, seed: int | None = None) -> dict:
if seed is not None:
random.seed(seed)
if task_id is None:
task_id = random.choice(list(self.tasks.keys()))
if task_id not in self.tasks:
raise KeyError(f"Unknown task_id: {task_id!r}")
task = self.tasks[task_id]
episode_id = str(uuid.uuid4())
self.episodes[episode_id] = {
"episode_id": episode_id,
"task": task,
"step_count": 0,
"done": False,
"rewards": [],
"last_test_results": [],
}
observation = self._initial_observation(task)
return {"episode_id": episode_id, "observation": observation}
def step(self, episode_id: str, action: dict) -> dict:
if episode_id not in self.episodes:
raise KeyError(f"Unknown episode_id: {episode_id!r}")
ep = self.episodes[episode_id]
if ep["done"]:
raise ValueError("Episode is already finished. Call reset() to start a new episode.")
task = ep["task"]
submitted_code = action.get("code", "")
ep["step_count"] += 1
test_results_raw, stdout, stderr = run_code_safely(
submitted_code,
task["test_code"],
timeout=10,
)
tests_passed = sum(1 for t in test_results_raw if t.get("passed", False))
total_tests = len(test_results_raw)
reward = round(tests_passed / total_tests, 4) if total_tests > 0 else 0.0
max_steps = task.get("max_steps", 5)
done = reward == 1.0 or ep["step_count"] >= max_steps
ep["done"] = done
ep["rewards"].append(reward)
ep["last_test_results"] = test_results_raw
observation = {
"task_id": task["task_id"],
"difficulty": task["difficulty"],
"description": task["description"],
"buggy_code": task["buggy_code"],
"test_descriptions": task["test_descriptions"],
"test_results": test_results_raw,
"stdout": stdout,
"stderr": stderr,
"step_count": ep["step_count"],
"max_steps": max_steps,
"reward": reward,
"done": done,
"total_tests": total_tests,
"tests_passed": tests_passed,
}
return {"observation": observation, "reward": reward, "done": done, "info": {}}
def state(self, episode_id: str) -> dict:
if episode_id not in self.episodes:
raise KeyError(f"Unknown episode_id: {episode_id!r}")
ep = self.episodes[episode_id]
task = ep["task"]
last_results = ep.get("last_test_results", [])
return {
"episode_id": episode_id,
"task_id": task["task_id"],
"difficulty": task["difficulty"],
"step_count": ep["step_count"],
"max_steps": task.get("max_steps", 5),
"last_reward": ep["rewards"][-1] if ep["rewards"] else 0.0,
"cumulative_reward": round(sum(ep["rewards"]), 4),
"tests_passed": sum(1 for t in last_results if t.get("passed", False)),
"total_tests": len(last_results),
"done": ep["done"],
}
def list_tasks(self) -> list[dict]:
return [
{
"task_id": t["task_id"],
"difficulty": t["difficulty"],
"description": t["description"],
"max_steps": t.get("max_steps", 5),
"total_tests": len(t["test_descriptions"]),
}
for t in self.tasks.values()
]
# ------------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------------
def _initial_observation(self, task: dict) -> dict:
return {
"task_id": task["task_id"],
"difficulty": task["difficulty"],
"description": task["description"],
"buggy_code": task["buggy_code"],
"test_descriptions": task["test_descriptions"],
"test_results": [],
"stdout": "",
"stderr": "",
"step_count": 0,
"max_steps": task.get("max_steps", 5),
"reward": 0.0,
"done": False,
"total_tests": len(task["test_descriptions"]),
"tests_passed": 0,
}
|