File size: 6,251 Bytes
afedf8c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
"""
Bug Report Structuring Environment - Core Environment Logic

Implements the OpenEnv 3-method interface:
  - reset(task_id) โ†’ initial observation with messy bug report
  - step(action) โ†’ graded observation with score + feedback
  - state() โ†’ episode metadata

The environment challenges an LLM agent to convert messy, unstructured
bug reports into well-organized structured formats.
"""

import uuid
import random
from typing import Optional

from models import (
    BugReportAction,
    BugReportObservation,
    BugReportState,
)
from tasks import get_task, get_all_task_ids
from graders import grade_submission


class BugReportEnvironment:
    """
    OpenEnv-compatible environment for bug report structuring.

    Supports concurrent sessions via per-instance state.
    """

    SUPPORTS_CONCURRENT_SESSIONS = True

    def __init__(self):
        self._state = BugReportState()
        self._task = None
        self._done = False
        self._best_score = 0.0
        self._rewards = []
        self._raw_report = ""

    def reset(
        self,
        task_id: Optional[str] = None,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
    ) -> BugReportObservation:
        """
        Initialize a new episode.

        Args:
            task_id: 'easy', 'medium', or 'hard'. Random if None.
            seed: Random seed for reproducibility.
            episode_id: Custom episode ID.

        Returns:
            Initial observation containing the messy bug report.
        """
        # Set seed for reproducibility
        if seed is not None:
            random.seed(seed)

        # Select task
        if task_id is None:
            task_id = random.choice(get_all_task_ids())

        self._task = get_task(task_id)
        self._done = False
        self._best_score = 0.0
        self._rewards = []
        self._raw_report = self._task["raw_report"]

        # Initialize state
        eid = episode_id or str(uuid.uuid4())
        self._state = BugReportState(
            episode_id=eid,
            step_count=0,
            task_id=task_id,
            max_steps=self._task["max_steps"],
            current_score=0.0,
            best_score=0.0,
            done=False,
            rewards=[],
        )

        return BugReportObservation(
            raw_report=self._raw_report,
            feedback=(
                f"New episode started. Task: {task_id} "
                f"(max {self._task['max_steps']} steps).\n"
                f"Read the messy bug report below and submit a structured version.\n\n"
                f"Required fields: title, steps_to_reproduce, expected_behavior, "
                f"actual_behavior, severity (low/medium/high/critical), environment"
            ),
            score=0.0,
            field_scores={},
            done=False,
            reward=0.0,
            step_count=0,
            task_id=task_id,
            max_steps=self._task["max_steps"],
        )

    def step(self, action: BugReportAction) -> BugReportObservation:
        """
        Process an agent's structured bug report submission.

        Args:
            action: BugReportAction with structured fields.

        Returns:
            Observation with score, feedback, and done status.
        """
        # Check if episode is active
        if self._task is None:
            return BugReportObservation(
                raw_report="",
                feedback="Error: No active episode. Call reset() first.",
                score=0.0,
                field_scores={},
                done=True,
                reward=0.0,
                step_count=0,
                task_id="",
                max_steps=0,
            )

        if self._done:
            return BugReportObservation(
                raw_report=self._raw_report,
                feedback=(
                    f"Episode already completed. Best score: {self._best_score:.2f}. "
                    f"Call reset() for a new episode."
                ),
                score=self._best_score,
                field_scores={},
                done=True,
                reward=0.0,
                step_count=self._state.step_count,
                task_id=self._state.task_id,
                max_steps=self._state.max_steps,
            )

        # Increment step
        self._state.step_count += 1

        # Grade the submission
        action_dict = action.model_dump()
        score, field_scores, feedback = grade_submission(action_dict, self._task)

        # Track best score and rewards
        self._best_score = max(self._best_score, score)

        # Reward = improvement over previous best (reward shaping)
        prev_best = self._state.best_score
        reward = max(0.0, score - prev_best)  # Only reward improvement
        if self._state.step_count == 1:
            reward = score  # First step gets full score as reward

        self._rewards.append(round(reward, 2))

        # Check if done
        at_max_steps = self._state.step_count >= self._state.max_steps
        high_score = score >= 0.95  # Near-perfect score ends early
        self._done = at_max_steps or high_score

        # Update state
        self._state.current_score = score
        self._state.best_score = self._best_score
        self._state.done = self._done
        self._state.rewards = self._rewards.copy()

        # Add done info to feedback
        if self._done:
            if high_score:
                feedback += f"\n\n๐ŸŽ‰ Excellent! Score {score:.2f} achieved. Episode complete."
            else:
                feedback += (
                    f"\n\nMax steps reached. Best score: {self._best_score:.2f}. "
                    f"Episode complete."
                )

        return BugReportObservation(
            raw_report=self._raw_report,
            feedback=feedback,
            score=score,
            field_scores=field_scores,
            done=self._done,
            reward=round(reward, 2),
            step_count=self._state.step_count,
            task_id=self._state.task_id,
            max_steps=self._state.max_steps,
        )

    @property
    def state(self) -> BugReportState:
        """Return current episode state."""
        return self._state