File size: 6,115 Bytes
b43aff5 3ca9f73 b43aff5 991b3a8 b43aff5 c7b11e5 b8c69e2 c7b11e5 b8c69e2 c7b11e5 b8c69e2 b43aff5 c7b11e5 b43aff5 c7b11e5 b43aff5 b1c5449 b43aff5 c7b11e5 | 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 158 159 160 161 162 163 164 | """ConfigDebugEnvironment - OpenEnv-compatible environment class.
Inherits from openenv.core.env_server.Environment and implements
the standard reset/step/state interface with multi-task logic.
"""
from typing import Optional, Any
from uuid import uuid4
from openenv.core.env_server import Environment
from server.models import ConfigDebugAction, ConfigDebugObservation, ConfigDebugState
from server.tasks.task_registry import get_task, TASK_ORDER
MAX_STEPS_PER_TASK = 5
class ConfigDebugEnvironment(Environment):
"""Multi-task config debugging environment.
Manages 7 sequential tasks internally. Each WebSocket session
(via create_fastapi_app) gets its own instance with independent state.
"""
SUPPORTS_CONCURRENT_SESSIONS = True
def __init__(self):
self._init_episode()
def _init_episode(self):
self.task_ids = list(TASK_ORDER)
self.current_task_index = 0
self.current_step = 0
self.total_reward = 0.01
self._done = False
self.tasks_completed: list = []
self.bugs_found_so_far = 0
self.previous_reward = 0.01
self.current_error_message: Optional[str] = None
self.current_broken_config: Optional[str] = None
self._episode_id = str(uuid4())
self._global_step = 0
# ---- OpenEnv interface methods ----
def reset(self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs: Any) -> ConfigDebugObservation:
"""Reset environment to initial state (task 1)."""
self._init_episode()
if episode_id:
self._episode_id = episode_id
return self._build_observation()
def step(self, action: ConfigDebugAction, timeout_s: Optional[float] = None, **kwargs: Any) -> ConfigDebugObservation:
"""Process an action: run the grader, advance tasks if done."""
if self._done:
return self._build_observation()
task_id = self._current_task_id()
task = get_task(task_id)
# Run the grader - returns (reward, error_message, bugs_fixed) tuple
grader_result = task.grader(action.fixed_config)
# Parse grader result
if isinstance(grader_result, tuple) and len(grader_result) >= 3:
reward = float(grader_result[0])
error_message = str(grader_result[1])
bugs_fixed = list(grader_result[2])
elif isinstance(grader_result, tuple) and len(grader_result) >= 1:
reward = float(grader_result[0])
error_message = ""
bugs_fixed = []
elif isinstance(grader_result, (int, float)):
reward = float(grader_result)
error_message = ""
bugs_fixed = []
else:
reward = 0.01
error_message = "Grader returned unexpected format"
bugs_fixed = []
reward = max(0.01, min(0.99, reward))
self.current_step += 1
self._global_step += 1
self.bugs_found_so_far = len(bugs_fixed)
self.previous_reward = round(reward, 4)
self.current_error_message = error_message
# Check if task is complete
task_done = reward >= 0.99 or self.current_step >= MAX_STEPS_PER_TASK
if task_done:
self.total_reward += reward
self.tasks_completed.append(task_id)
self.current_task_index += 1
self.current_step = 0
self.bugs_found_so_far = 0
self.current_error_message = None
self.current_broken_config = None
if self.current_task_index >= len(self.task_ids):
self._done = True
else:
self.current_broken_config = action.fixed_config
obs = self._build_observation()
obs.done = self._done
obs.reward = round(reward, 4)
return obs
@property
def state(self) -> ConfigDebugState:
"""Return current environment state with enhanced RL signals."""
tasks_remaining = self.task_ids[self.current_task_index:]
if self._done:
tasks_remaining = []
total_tasks = len(self.task_ids)
completed_tasks = len(self.tasks_completed)
progress_ratio = completed_tasks / total_tasks if total_tasks > 0 else 0.01
current_task = get_task(self._current_task_id())
return ConfigDebugState(
episode_id=self._episode_id,
step_count=self._global_step,
current_task_id=self._current_task_id(),
current_step=self.current_step,
max_steps=MAX_STEPS_PER_TASK,
total_reward=round(self.total_reward, 4),
is_done=self._done,
tasks_completed=list(self.tasks_completed),
tasks_remaining=tasks_remaining,
bugs_found_so_far=self.bugs_found_so_far,
current_error_message=self.current_error_message,
progress_ratio=round(progress_ratio, 2),
current_difficulty=current_task.difficulty,
)
# ---- Internal helpers ----
def _current_task_id(self) -> str:
if self.current_task_index < len(self.task_ids):
return self.task_ids[self.current_task_index]
return self.task_ids[-1]
def _build_observation(self) -> ConfigDebugObservation:
task_id = self._current_task_id()
task = get_task(task_id)
broken = self.current_broken_config if self.current_broken_config is not None else task.broken_config
error = self.current_error_message if self.current_error_message is not None else task.error_message
return ConfigDebugObservation(
broken_config=broken,
ground_truth=task.ground_truth,
file_type=task.file_type,
error_message=error,
task_id=task.task_id,
task_description=task.description,
difficulty=task.difficulty,
num_bugs=task.num_bugs,
bugs_found_so_far=self.bugs_found_so_far,
previous_reward=self.previous_reward,
done=self._done,
reward=self.previous_reward,
) |