File size: 5,138 Bytes
71a1c53
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""

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