Spaces:
Sleeping
Sleeping
File size: 8,902 Bytes
b92d20c 2cd9867 d0cfac3 2cd9867 d0cfac3 b92d20c 6139c83 b92d20c 2d2c6e4 b92d20c 2cd9867 b92d20c 2cd9867 b92d20c 2cd9867 b92d20c 2d2c6e4 b92d20c 2d2c6e4 b92d20c ee45044 b92d20c d0cfac3 b92d20c 2d2c6e4 b92d20c 2d2c6e4 b92d20c 2d2c6e4 b92d20c 6139c83 b92d20c 2d2c6e4 b92d20c d0cfac3 b92d20c d0cfac3 b92d20c d0cfac3 b92d20c | 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 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | from __future__ import annotations
from typing import Any, Dict, List, Optional, Set
from openenv.core import Environment
from codereview_env.models import (
CodeReviewAction,
CodeReviewObservation,
CodeReviewState,
ReviewArtifact,
ReviewFinding,
)
from server.reward import RewardBreakdown, RewardComputer
from server.tasks import TASKS, TASKS_BY_ID, ReviewTask
_MIN_PUBLIC_SCORE = 0.05
_MAX_PUBLIC_SCORE = 0.95
def _clamp_score(raw: float) -> float:
"""Keep validator-visible scores inside a conservative open interval."""
return max(_MIN_PUBLIC_SCORE, min(_MAX_PUBLIC_SCORE, raw))
class CodeReviewEnvironment(
Environment[CodeReviewObservation, CodeReviewAction, CodeReviewState]
):
SUPPORTS_CONCURRENT_SESSIONS = True
def __init__(self, reward_computer: Optional[RewardComputer] = None):
super().__init__()
self.reward_computer = reward_computer or RewardComputer()
self.task: Optional[ReviewTask] = None
self._opened_artifact_ids: Set[str] = set()
self._submitted_findings: List[ReviewFinding] = []
self._recent_events: List[str] = []
self._cumulative_reward = _MIN_PUBLIC_SCORE
self._score = _MIN_PUBLIC_SCORE
self._last_action_error: Optional[str] = None
self._step_limit = 0
self._task_index = 0
self._episode_done = False
self.episode_id = None
self.step_count = 0
def get_metadata(self) -> Dict[str, Any]:
return {
"name": "codereview-env",
"domain": "software engineering",
"description": "A deterministic code review benchmark for real pull request triage.",
"tasks": [task.task_id for task in TASKS],
}
def reset(
self,
seed: Optional[int] = None,
episode_id: Optional[str] = None,
**kwargs: Any,
) -> CodeReviewObservation:
task_id = kwargs.get("task_id")
if isinstance(seed, str) and not task_id:
task_id = seed
seed = None
if task_id:
task = TASKS_BY_ID[task_id]
elif seed is not None:
task = TASKS[seed % len(TASKS)]
else:
task = TASKS[self._task_index % len(TASKS)]
self._task_index += 1
self.task = task
self.episode_id = episode_id or task.task_id
self.step_count = 0
self._step_limit = task.step_limit
self._opened_artifact_ids = set(task.starting_artifacts)
self._submitted_findings = []
self._recent_events = ["Episode reset."]
self._cumulative_reward = _MIN_PUBLIC_SCORE
self._score = _MIN_PUBLIC_SCORE
self._last_action_error = None
self._episode_done = False
return self._build_observation(reward=_MIN_PUBLIC_SCORE, done=False)
def step(
self, action: CodeReviewAction, timeout_s: Optional[float] = None, **kwargs: Any
) -> CodeReviewObservation:
if self.task is None:
raise RuntimeError("reset() must be called before step().")
if self._episode_done:
self._last_action_error = "Episode already finished."
return self._build_observation(reward=_MIN_PUBLIC_SCORE, done=True)
self.step_count += 1
reward_breakdown: RewardBreakdown
if action.action_type == "open_artifact":
reward_breakdown = self._handle_open_artifact(action)
done = False
elif action.action_type == "submit_review":
reward_breakdown = self._handle_submit(action)
done = True
else:
reward_breakdown = self.reward_computer.invalid_action(
"Unsupported action_type."
)
done = False
if self.step_count >= self._step_limit and not done:
self._recent_events.append(
"Step limit reached before review was submitted."
)
done = True
self._episode_done = done
self._last_action_error = reward_breakdown.last_action_error
self._cumulative_reward = max(0.05, min(0.95, self._cumulative_reward + reward_breakdown.reward))
self._recent_events.append(
f"Step {self.step_count}: {action.action_type} -> reward {reward_breakdown.reward:.2f}"
)
return self._build_observation(
reward=reward_breakdown.reward,
done=done,
extra_metadata={
"reward_breakdown": reward_breakdown.components,
"grader_details": reward_breakdown.grader_details,
},
)
@property
def state(self) -> CodeReviewState:
task_metadata = {
"step_limit": self._step_limit,
"done": self._episode_done,
}
return CodeReviewState(
episode_id=self.episode_id,
step_count=self.step_count,
task_id=self.task.task_id if self.task else None,
difficulty=self.task.difficulty if self.task else None,
title=self.task.title if self.task else None,
opened_artifact_ids=sorted(self._opened_artifact_ids),
submitted_findings=self._submitted_findings,
cumulative_reward=round(_clamp_score(self._cumulative_reward), 6),
score=round(_clamp_score(self._score), 6),
last_action_error=self._last_action_error,
task_metadata=task_metadata,
)
def _handle_open_artifact(self, action: CodeReviewAction) -> RewardBreakdown:
artifact_id = action.artifact_id
if not artifact_id:
return self.reward_computer.invalid_action(
"artifact_id is required for open_artifact."
)
if artifact_id not in self.task.artifacts:
return self.reward_computer.invalid_action(
f"Unknown artifact_id: {artifact_id}"
)
repeated = artifact_id in self._opened_artifact_ids
self._opened_artifact_ids.add(artifact_id)
artifact = self.task.artifacts[artifact_id]
self._recent_events.append(f"Opened {artifact.title}.")
return self.reward_computer.artifact_reward(
self.task, artifact_id, self._opened_artifact_ids, repeated
)
def _handle_submit(self, action: CodeReviewAction) -> RewardBreakdown:
self._submitted_findings = list(action.findings)
breakdown = self.reward_computer.submission_reward(
self.task,
action.findings,
self._opened_artifact_ids,
self.step_count,
self._step_limit,
)
self._score = _clamp_score(breakdown.score)
self._recent_events.append(f"Submitted {len(action.findings)} findings.")
return breakdown
def _build_observation(
self,
reward: float,
done: bool,
extra_metadata: Optional[Dict[str, Any]] = None,
) -> CodeReviewObservation:
opened_artifacts = []
available_artifacts = []
if self.task is not None:
for artifact_id, artifact in self.task.artifacts.items():
model = ReviewArtifact(
artifact_id=artifact.artifact_id,
kind=artifact.kind, # type: ignore[arg-type]
title=artifact.title,
preview=artifact.preview,
opened=artifact_id in self._opened_artifact_ids,
content=(
artifact.content
if artifact_id in self._opened_artifact_ids
else None
),
)
available_artifacts.append(model)
if model.opened:
opened_artifacts.append(model)
metadata = {
"task_id": self.task.task_id if self.task else None,
"score": round(_clamp_score(self._score), 6),
"step_count": self.step_count,
"opened_artifact_ids": sorted(self._opened_artifact_ids),
}
if extra_metadata:
metadata.update(extra_metadata)
return CodeReviewObservation(
task_id=self.task.task_id if self.task else "",
difficulty=self.task.difficulty if self.task else "easy",
title=self.task.title if self.task else "",
objective=self.task.objective if self.task else "",
summary=self.task.summary if self.task else "",
step_limit=self._step_limit,
opened_artifacts=opened_artifacts,
available_artifacts=available_artifacts,
recent_events=self._recent_events[-6:],
last_action_error=self._last_action_error,
score=round(_clamp_score(self._score), 6),
done=done,
reward=round(_clamp_score(reward), 6),
metadata=metadata,
)
|