commit

environment/env.py

@@ -0,0 +1,135 @@
+"""
+CodeReviewEnv — main environment logic.
+Manages state, episode flow, reward accumulation, and per-step grading.
+"""
+from typing import Optional, List, Dict, Any
+
+from environment.models import (
+    CodeReviewAction,
+    CodeReviewObservation,
+    StepResult,
+    ResetResult,
+    StateResult,
+)
+from environment.tasks import get_task, list_tasks as _list_tasks
+from environment.graders import grade
+
+
+class CodeReviewEnv:
+    """
+    OpenEnv-compliant code-review environment.
+
+    Episode flow
+    ────────────
+    reset(task_id) → observation
+    step(action)   → (observation, reward, done, info)   [repeated ≤ max_steps]
+    state()        → current state snapshot
+    """
+
+    def __init__(self) -> None:
+        self._task: Optional[dict] = None
+        self._step_number: int = 0
+        self._total_reward: float = 0.0
+        self._actions_history: List[Dict[str, Any]] = []
+        self._done: bool = False
+        self._initialized: bool = False
+        self._last_feedback: Optional[str] = None
+
+    # ── public properties ────────────────────────────────────────────────────
+
+    @property
+    def is_initialized(self) -> bool:
+        return self._initialized
+
+    # ── core API ─────────────────────────────────────────────────────────────
+
+    def reset(self, task_id: Optional[str] = None) -> ResetResult:
+        """Start a new episode. Defaults to the easy task."""
+        if task_id is None:
+            task_id = "easy_syntax"
+
+        self._task = get_task(task_id)
+        self._step_number = 0
+        self._total_reward = 0.0
+        self._actions_history = []
+        self._done = False
+        self._initialized = True
+        self._last_feedback = None
+
+        obs = self._make_observation()
+        return ResetResult(observation=obs)
+
+    def step(self, action: CodeReviewAction) -> StepResult:
+        """Process one agent action and return (observation, reward, done, info)."""
+        if not self._initialized or self._done:
+            raise RuntimeError("Call reset() before stepping, or episode is over.")
+
+        self._step_number += 1
+        task_id = self._task["task_id"]
+        ground_truth = self._task["ground_truth"]
+
+        # Grade the action
+        reward, feedback = grade(task_id, action, ground_truth)
+        self._last_feedback = feedback
+        self._total_reward += reward
+
+        # Record history
+        self._actions_history.append(
+            {
+                "step": self._step_number,
+                "num_issues_reported": len(action.identified_issues),
+                "has_fix": action.suggested_fix is not None,
+                "reward": reward,
+            }
+        )
+
+        # Episode ends when: agent says done, reward is perfect, or max steps reached
+        max_steps = self._task["max_steps"]
+        done = action.done or reward >= 0.95 or self._step_number >= max_steps
+        self._done = done
+
+        obs = self._make_observation()
+
+        return StepResult(
+            observation=obs,
+            reward=reward,
+            done=done,
+            info={
+                "feedback": feedback,
+                "step": self._step_number,
+                "total_reward": round(self._total_reward, 4),
+                "cumulative_score": round(
+                    self._total_reward / max(self._step_number, 1), 4
+                ),
+            },
+        )
+
+    def get_state(self) -> StateResult:
+        """Return a snapshot of the current episode state."""
+        return StateResult(
+            task_id=self._task["task_id"] if self._task else "",
+            step_number=self._step_number,
+            total_reward=round(self._total_reward, 4),
+            actions_history=self._actions_history,
+            done=self._done,
+            initialized=self._initialized,
+        )
+
+    def list_tasks(self) -> list:
+        return _list_tasks()
+
+    # ── internal helpers ─────────────────────────────────────────────────────
+
+    def _make_observation(self) -> CodeReviewObservation:
+        t = self._task
+        return CodeReviewObservation(
+            task_id=t["task_id"],
+            task_name=t["task_name"],
+            difficulty=t["difficulty"],
+            language=t["language"],
+            code_snippet=t["code_snippet"],
+            context=t["context"],
+            step_number=self._step_number,
+            max_steps=t["max_steps"],
+            previous_feedback=self._last_feedback,
+        )