code-debug-env / code_debug_env /server /debug_environment.py
shivammmmm's picture
Initial commit: code debug environment
71a1c53
Raw
History Blame Contribute Delete
5.14 kB
"""
Code Debug Environment (server side)
-------------------------------------
Presents buggy Python functions and grades agent-submitted fixes.
"""
from __future__ import annotations
import itertools
import os
import random
import threading
import uuid
from openenv.core.env_server.interfaces import Action, Environment, Observation
from ..models import DebugAction, DebugObservation, DebugState
from .grader import grade
from .problems import ALL_TASK_NAMES, TASK_PROBLEMS, Problem
MAX_STEPS = int(os.getenv("CODE_DEBUG_MAX_STEPS", "10"))
COMPLETION_BONUS = 0.3
SYNTAX_PENALTY = -0.1
NO_CHANGE_PENALTY = -0.05
STEP_COST = -0.01
_task_cycle = itertools.cycle(ALL_TASK_NAMES)
_task_cycle_lock = threading.Lock()
def _next_task() -> str:
with _task_cycle_lock:
return next(_task_cycle)
class CodeDebugEnvironment(Environment):
"""
RL environment for iterative code debugging.
On reset the environment picks a problem for the requested task,
runs the buggy code against its test suite, and returns the initial
observation. Each step accepts a DebugAction containing the agent's
corrected code, grades it, and returns an updated observation with
reward.
Task selection:
- If ``CODE_DEBUG_TASK`` env var is set, that task is always used.
- Otherwise tasks cycle: fix_syntax -> fix_logic -> fix_algorithm.
"""
def __init__(self) -> None:
self._state = DebugState()
self._problem: Problem | None = None
self._prev_code: str = ""
self._prev_passed: int = 0
# ------------------------------------------------------------------
# OpenEnv interface
# ------------------------------------------------------------------
def reset(self) -> Observation:
task_override = os.getenv("CODE_DEBUG_TASK")
task_name = task_override if task_override else _next_task()
problems = TASK_PROBLEMS.get(task_name)
if not problems:
task_name = "fix_syntax"
problems = TASK_PROBLEMS[task_name]
self._problem = random.choice(problems)
initial = grade(
self._problem.buggy_code,
self._problem.function_name,
self._problem.tests,
)
self._state = DebugState(
episode_id=str(uuid.uuid4()),
step_count=0,
current_task=task_name,
best_tests_passed=initial.tests_passed,
tests_total=initial.tests_total,
)
self._prev_code = self._problem.buggy_code
self._prev_passed = initial.tests_passed
return DebugObservation(
task_name=task_name,
buggy_code=self._problem.buggy_code,
task_description=self._problem.description,
current_code=self._problem.buggy_code,
test_results=initial.details,
tests_passed=initial.tests_passed,
tests_total=initial.tests_total,
stderr=initial.stderr,
)
def step(self, action: Action) -> Observation:
if not isinstance(action, DebugAction):
raise ValueError(f"Expected DebugAction, got {type(action)}")
if self._problem is None:
raise RuntimeError("Call reset() before step()")
self._state.step_count += 1
submitted_code = action.code
result = grade(
submitted_code,
self._problem.function_name,
self._problem.tests,
)
# --- reward computation ---
reward = STEP_COST
if result.stderr and result.tests_passed == 0 and "SyntaxError" in result.stderr:
reward += SYNTAX_PENALTY
elif submitted_code.strip() == self._prev_code.strip():
reward += NO_CHANGE_PENALTY
else:
improvement = result.tests_passed - self._prev_passed
if result.tests_total > 0:
reward += improvement / result.tests_total
all_passed = result.tests_passed == result.tests_total
if all_passed:
reward += COMPLETION_BONUS
done = all_passed or self._state.step_count >= MAX_STEPS
# track best
if result.tests_passed > self._state.best_tests_passed:
self._state.best_tests_passed = result.tests_passed
self._state.tests_total = result.tests_total
self._prev_code = submitted_code
self._prev_passed = result.tests_passed
obs = DebugObservation(
task_name=self._state.current_task,
buggy_code=self._problem.buggy_code,
task_description=self._problem.description,
current_code=submitted_code,
test_results=result.details,
tests_passed=result.tests_passed,
tests_total=result.tests_total,
stderr=result.stderr,
reward=reward,
done=done,
)
return obs
@property
def state(self) -> DebugState:
return self._state