Mirror of GitHub source: OpenEnv-compliant LeniencyBench environment + training scripts
6b4f87f verified | """DriftEnv — OpenEnv-compliant multi-step environment for the | |
| Policy-Drift support-triage task. | |
| """ | |
| from __future__ import annotations | |
| from typing import Optional, List, Set | |
| from drift_env.models import Action, Email, EmailKind, Observation, State, StepResult | |
| from drift_env.episodes import Episode, generate_episode | |
| from drift_env.grader import grade_step | |
| def _email_to_summary(email: Email, action_taken: Optional[Action] = None) -> dict: | |
| entry = { | |
| "email_id": email.id, | |
| "kind": email.kind.value, | |
| "subject": email.subject, | |
| "body": email.body, | |
| "sender": email.sender, | |
| } | |
| if action_taken is not None: | |
| entry["action_taken"] = action_taken.action_type.value | |
| return entry | |
| def _email_for_observation(email: Email) -> Email: | |
| """Strip grader-only metadata before exposing to agent.""" | |
| return Email( | |
| id=email.id, | |
| kind=email.kind, | |
| subject=email.subject, | |
| body=email.body, | |
| sender=email.sender, | |
| meta={}, | |
| ) | |
| class DriftEnv: | |
| def __init__(self) -> None: | |
| self._episode: Optional[Episode] = None | |
| self._index: int = 0 | |
| self._cumulative: float = 0.0 | |
| self._done: bool = True | |
| self._history: List[dict] = [] | |
| self._armed_drifts: Set[str] = set() | |
| # ------------------------------------------------------------------ | |
| # OpenEnv API | |
| # ------------------------------------------------------------------ | |
| def reset(self, seed: int = 0, episode_id: str = "ep_0") -> Observation: | |
| self._episode = generate_episode(seed=seed, episode_id=episode_id) | |
| self._index = 0 | |
| self._cumulative = 0.0 | |
| self._done = False | |
| self._history = [] | |
| self._armed_drifts = set() | |
| return self._current_observation() | |
| def step(self, action: Action) -> StepResult: | |
| if self._episode is None: | |
| raise RuntimeError("Call reset() before step().") | |
| if self._done: | |
| raise RuntimeError("Episode is done. Call reset() to start a new one.") | |
| step_record = self._episode.steps[self._index] | |
| is_admin = step_record.email.kind == EmailKind.ADMIN | |
| reward, breakdown, drift_to_clear = grade_step( | |
| action=action, | |
| hint=step_record.correct_action_hint, | |
| email_meta=step_record.email.meta, | |
| drift_sensitive_to=step_record.drift_sensitive_to, | |
| armed_drifts=self._armed_drifts, | |
| is_admin_email=is_admin, | |
| ) | |
| # If agent proved awareness of a drift, clear the arm. | |
| if drift_to_clear: | |
| self._armed_drifts.discard(drift_to_clear) | |
| # If THIS step is an admin email, arm its drift for future steps. | |
| if is_admin: | |
| drift_name = step_record.email.meta.get("drift_event") | |
| if drift_name: | |
| self._armed_drifts.add(drift_name) | |
| self._cumulative += reward | |
| # Record this email + the action for history. | |
| self._history.append(_email_to_summary(step_record.email, action_taken=action)) | |
| self._index += 1 | |
| self._done = self._index >= len(self._episode.steps) | |
| next_obs = None if self._done else self._current_observation() | |
| info = { | |
| "episode_id": self._episode.id, | |
| "step": self._index, | |
| "total_steps": len(self._episode.steps), | |
| "correct_action_hint": step_record.correct_action_hint, | |
| "drift_sensitive_to": step_record.drift_sensitive_to, | |
| "armed_drifts_after": sorted(self._armed_drifts), | |
| "breakdown": breakdown, | |
| "cumulative_reward": round(self._cumulative, 4), | |
| } | |
| return StepResult( | |
| observation=next_obs, | |
| reward=round(reward, 4), | |
| done=self._done, | |
| info=info, | |
| ) | |
| def state(self) -> State: | |
| return State( | |
| episode_id=self._episode.id if self._episode else None, | |
| email_index=self._index, | |
| total_emails=len(self._episode.steps) if self._episode else 0, | |
| done=self._done, | |
| cumulative_reward=round(self._cumulative, 4), | |
| ) | |
| # ------------------------------------------------------------------ | |
| def _current_observation(self) -> Observation: | |
| assert self._episode is not None | |
| step = self._episode.steps[self._index] | |
| return Observation( | |
| current_email=_email_for_observation(step.email), | |
| email_index=self._index, | |
| total_emails=len(self._episode.steps), | |
| inbox_history=list(self._history), | |
| ) | |