feat: rewrite env to be fully openenv-core compliant

environment/actions.py

@@ -1,6 +1,7 @@
 from enum import Enum
 from typing import Optional
 from pydantic import BaseModel, field_validator
+from openenv.core import Action, Observation, State
 
 
 class ActionType(str, Enum):
@@ -10,7 +11,7 @@ class ActionType(str, Enum):
     submit_answer = "submit_answer"
 
 
-class ContextCorruptionAction(BaseModel):
+class ContextCorruptionAction(Action):
     action_type: ActionType
     doc_id: Optional[int] = None
     answer: Optional[str] = None
@@ -31,12 +32,22 @@ class Document(BaseModel):
     is_flagged: bool = False
 
 
-class EpisodeObservation(BaseModel):
-    question: str
-    documents: list[Document]
-    flagged_ids: list[int]
-    budget_remaining: int
-    turn: int
-    episode_done: bool = False
-    reward: Optional[float] = None
+class EpisodeObservation(Observation):
+    question: str = ""
+    documents: list[Document] = []
+    flagged_ids: list[int] = []
+    budget_remaining: int = 0
+    turn: int = 0
     message: Optional[str] = None
+    # `done` and `reward` inherited from Observation
+
+
+class ContextCorruptionState(State):
+    question: str = ""
+    ground_truth: str = ""
+    corrupt_ids: list[int] = []
+    flagged_ids: list[int] = []
+    budget_used: int = 0
+    done: bool = False
+    reward: Optional[float] = None
+    breakdown: Optional[dict] = None

environment/env.py

@@ -2,20 +2,28 @@ import json
 import random
 from pathlib import Path
 
-from environment.actions import ActionType, ContextCorruptionAction, Document, EpisodeObservation
-from environment.reward import compute_reward
+from openenv.core import Environment
+
+from environment.actions import (
+    ActionType, ContextCorruptionAction, Document,
+    EpisodeObservation, ContextCorruptionState,
+)
+from environment.reward import ContextCorruptionRubric
 
 _FALLBACK_FACTS = [
     {"question": "What is the capital of France?", "answer": "Paris"}
 ]
 
 
-class ContextCorruptionEnv:
+class ContextCorruptionEnv(Environment[ContextCorruptionAction, EpisodeObservation, ContextCorruptionState]):
     MAX_BUDGET = 12
     NUM_DOCS = 8
     DIFFICULTY_LEVELS = [1, 2, 3, 4]
+    SUPPORTS_CONCURRENT_SESSIONS = True
 
     def __init__(self, difficulty=None):
+        rubric = ContextCorruptionRubric(state_fn=self._state_dict)
+        super().__init__(rubric=rubric)
         self.difficulty = difficulty
         facts_path = Path(__file__).parent.parent / "data" / "facts.json"
         if facts_path.exists():
@@ -37,8 +45,11 @@ class ContextCorruptionEnv:
         self._reward = None
         self._breakdown = None
 
-    def reset(self) -> EpisodeObservation:
+    def reset(self, seed=None, episode_id=None, **kwargs) -> EpisodeObservation:
+        self._reset_rubric()
         self._reset_state()
+        if seed is not None:
+            random.seed(seed)
         fact = random.choice(self._facts)
         n_corrupt = self.difficulty if self.difficulty is not None else random.choice(self.DIFFICULTY_LEVELS)
         self._corrupt_ids = random.sample(range(self.NUM_DOCS), n_corrupt)
@@ -55,18 +66,17 @@ class ContextCorruptionEnv:
             ]
 
         self._documents = raw_docs
-        return self._build_observation()
+        return self._apply_transform(self._build_observation())
 
-    def step(self, action: ContextCorruptionAction) -> EpisodeObservation:
+    def step(self, action: ContextCorruptionAction, timeout_s=None, **kwargs) -> EpisodeObservation:
         if self._done:
-            return self._build_observation(message="Episode already done.")
+            return self._apply_transform(self._build_observation(message="Episode already done."))
 
         self._turn += 1
         self._budget_used += 1
-        reward = None
 
         if action.action_type == ActionType.read_doc:
-            pass  # budget cost is the point; content already in observation
+            pass
 
         elif action.action_type == ActionType.flag_suspicious:
             if action.doc_id is not None and action.doc_id not in self._flagged_ids:
@@ -77,48 +87,44 @@ class ContextCorruptionEnv:
                 self._flagged_ids.remove(action.doc_id)
 
         elif action.action_type == ActionType.submit_answer:
-            reward, self._breakdown = compute_reward(
-                submitted_answer=action.answer or "",
-                ground_truth_answer=self._ground_truth,
-                flagged_ids=self._flagged_ids,
-                corrupt_ids=self._corrupt_ids,
-                confidence=action.confidence or 0.0,
-                budget_used=self._budget_used,
-                max_budget=self.MAX_BUDGET,
-            )
-            self._reward = reward
             self._done = True
 
         # Force-submit on budget exhaustion
         if self._budget_used >= self.MAX_BUDGET and not self._done:
-            reward, self._breakdown = compute_reward(
-                submitted_answer="",
-                ground_truth_answer=self._ground_truth,
-                flagged_ids=self._flagged_ids,
-                corrupt_ids=self._corrupt_ids,
-                confidence=0.0,
-                budget_used=self._budget_used,
-                max_budget=self.MAX_BUDGET,
-            )
-            self._reward = reward
             self._done = True
 
-        return self._build_observation(reward=reward)
+        obs = self._build_observation()
+
+        if obs.done:
+            obs.reward = self._apply_rubric(action, obs)
+            self._reward = obs.reward
+            self._breakdown = self.rubric.last_breakdown if self.rubric else None
+
+        return self._apply_transform(obs)
+
+    @property
+    def state(self) -> ContextCorruptionState:
+        return ContextCorruptionState(
+            question=self._question,
+            ground_truth=self._ground_truth,
+            corrupt_ids=list(self._corrupt_ids),
+            flagged_ids=list(self._flagged_ids),
+            budget_used=self._budget_used,
+            done=self._done,
+            reward=self._reward,
+            breakdown=self._breakdown,
+        )
 
-    def state(self) -> dict:
+    def _state_dict(self) -> dict:
         return {
-            "question": self._question,
             "ground_truth": self._ground_truth,
-            "corrupt_ids": self._corrupt_ids,
-            "flagged_ids": self._flagged_ids,
+            "flagged_ids": list(self._flagged_ids),
+            "corrupt_ids": list(self._corrupt_ids),
             "budget_used": self._budget_used,
-            "turn": self._turn,
-            "done": self._done,
-            "reward": self._reward,
-            "breakdown": self._breakdown,
+            "max_budget": self.MAX_BUDGET,
         }
 
-    def _build_observation(self, reward=None, message=None) -> EpisodeObservation:
+    def _build_observation(self, message=None) -> EpisodeObservation:
         docs = [
             Document(
                 id=d["id"],
@@ -134,7 +140,7 @@ class ContextCorruptionEnv:
             flagged_ids=list(self._flagged_ids),
             budget_remaining=self.MAX_BUDGET - self._budget_used,
             turn=self._turn,
-            episode_done=self._done,
-            reward=reward if reward is not None else self._reward,
+            done=self._done,
+            reward=self._reward,
             message=message,
         )

environment/reward.py

@@ -1,4 +1,5 @@
 import re
+from openenv.core.rubrics import Rubric
 
 
 def _normalize(text: str) -> str:
@@ -17,24 +18,19 @@ def compute_reward(
     budget_used: int,
     max_budget: int,
 ) -> tuple[float, dict]:
-    # Answer correctness
     correct = _normalize(submitted_answer) == _normalize(ground_truth_answer)
     answer_score = 0.4 if correct else 0.0
 
-    # Flag recall
     true_positives = [i for i in flagged_ids if i in corrupt_ids]
     recall = len(true_positives) / len(corrupt_ids) if corrupt_ids else 0.0
     recall_score = 0.3 * recall
 
-    # Precision (false positive penalty)
     false_positives = [i for i in flagged_ids if i not in corrupt_ids]
     precision_score = max(0.0, 0.2 - 0.1 * len(false_positives))
 
-    # Confidence calibration
     confidence = confidence or 0.0
     calibration_score = (0.1 * confidence) if correct else (-0.2 * confidence)
 
-    # Efficiency bonus
     efficiency_score = 0.05 * (1 - budget_used / max_budget)
 
     total = answer_score + recall_score + precision_score + calibration_score + efficiency_score
@@ -48,3 +44,32 @@ def compute_reward(
         "total": round(total, 4),
     }
     return round(total, 4), breakdown
+
+
+class ContextCorruptionRubric(Rubric):
+    """Scores a completed episode using compute_reward().
+
+    Requires a state_fn closure to access ground-truth env state that is
+    intentionally hidden from the agent's observation.
+    """
+
+    def __init__(self, state_fn):
+        super().__init__()
+        self._state_fn = state_fn
+        self.last_breakdown: dict = {}
+
+    def forward(self, action, observation) -> float:
+        if not observation.done:
+            return 0.0
+        s = self._state_fn()
+        reward, breakdown = compute_reward(
+            submitted_answer=getattr(action, "answer", None) or "",
+            ground_truth_answer=s["ground_truth"],
+            flagged_ids=s["flagged_ids"],
+            corrupt_ids=s["corrupt_ids"],
+            confidence=getattr(action, "confidence", None) or 0.0,
+            budget_used=s["budget_used"],
+            max_budget=s["max_budget"],
+        )
+        self.last_breakdown = breakdown
+        return reward

environment/server.py

@@ -1,56 +1,16 @@
-import uuid
-from fastapi import FastAPI, HTTPException
+from openenv.core import create_app
 import uvicorn
 
 from environment.actions import ContextCorruptionAction, EpisodeObservation
 from environment.env import ContextCorruptionEnv
 
-app = FastAPI(title="ContextCorruption-Env")
-
-# session_id -> env instance
-_sessions: dict[str, ContextCorruptionEnv] = {}
-_MAX_SESSIONS = 64
-
-
-@app.post("/reset", response_model=EpisodeObservation)
-def reset(session_id: str | None = None):
-    if session_id is None:
-        if len(_sessions) >= _MAX_SESSIONS:
-            raise HTTPException(status_code=503, detail="Max concurrent sessions reached")
-        session_id = str(uuid.uuid4())
-    if session_id not in _sessions:
-        _sessions[session_id] = ContextCorruptionEnv()
-    obs = _sessions[session_id].reset()
-    return obs
-
-
-@app.post("/step/{session_id}", response_model=EpisodeObservation)
-def step(session_id: str, action: ContextCorruptionAction):
-    if session_id not in _sessions:
-        raise HTTPException(status_code=404, detail="Session not found")
-    obs = _sessions[session_id].step(action)
-    if obs.episode_done:
-        del _sessions[session_id]
-    return obs
-
-
-@app.get("/state/{session_id}")
-def state(session_id: str):
-    if session_id not in _sessions:
-        raise HTTPException(status_code=404, detail="Session not found")
-    return _sessions[session_id].state()
-
-
-@app.delete("/session/{session_id}")
-def close_session(session_id: str):
-    _sessions.pop(session_id, None)
-    return {"status": "closed"}
-
-
-@app.get("/health")
-def health():
-    return {"status": "ok", "active_sessions": len(_sessions)}
-
+app = create_app(
+    env=lambda: ContextCorruptionEnv(),
+    action_cls=ContextCorruptionAction,
+    observation_cls=EpisodeObservation,
+    env_name="ContextCorruption-Env",
+    max_concurrent_envs=64,
+)
 
 if __name__ == "__main__":
     uvicorn.run("environment.server:app", host="0.0.0.0", port=8000, reload=False)

eval/baseline_eval.py

@@ -33,7 +33,7 @@ def run_baseline():
                 )
 
             obs = env.step(action)
-            done = obs.episode_done
+            done = obs.done
 
         rewards.append(obs.reward)
         if (ep + 1) % 10 == 0:

requirements.txt

@@ -1,4 +1,5 @@
 accelerate==1.13.0
+aiofile==3.9.0
 aiohappyeyeballs==2.6.1
 aiohttp==3.13.5
 aiosignal==1.4.0
@@ -6,59 +7,118 @@ annotated-doc==0.0.4
 annotated-types==0.7.0
 anyio==4.13.0
 attrs==26.1.0
+audioop-lts==0.2.2
+Authlib==1.7.0
+beartype==0.22.9
+brotli==1.2.0
+cachetools==7.0.6
+caio==0.9.25
 certifi==2026.4.22
+cffi==2.0.0
 charset-normalizer==3.4.7
 click==8.3.3
+cryptography==47.0.0
+cyclopts==4.11.0
 datasets==4.8.4
 dill==0.4.1
+distro==1.9.0
+dnspython==2.8.0
+docstring_parser==0.18.0
+docutils==0.22.4
+email-validator==2.3.0
+exceptiongroup==1.3.1
 Faker==40.15.0
 fastapi==0.136.1
+fastmcp==3.2.4
 filelock==3.29.0
 frozenlist==1.8.0
 fsspec==2026.2.0
 gitdb==4.0.12
 GitPython==3.1.47
+gradio==6.13.0
+gradio_client==2.5.0
+griffelib==2.0.2
+groovy==0.1.2
 h11==0.16.0
+hf-gradio==0.4.1
 hf-xet==1.4.3
 httpcore==1.0.9
 httpx==0.28.1
+httpx-sse==0.4.3
 huggingface_hub==1.12.0
 idna==3.13
+importlib_metadata==8.7.1
+jaraco.classes==3.4.0
+jaraco.context==6.1.2
+jaraco.functools==4.4.0
 Jinja2==3.1.6
+jiter==0.14.0
+joserfc==1.6.4
+jsonref==1.1.0
+jsonschema==4.26.0
+jsonschema-path==0.4.5
+jsonschema-specifications==2025.9.1
+keyring==25.7.0
 markdown-it-py==4.0.0
 MarkupSafe==3.0.3
+mcp==1.27.0
 mdurl==0.1.2
+more-itertools==11.0.2
 mpmath==1.3.0
 multidict==6.7.1
 multiprocess==0.70.19
 networkx==3.6.1
 numpy==2.4.4
-openenv==0.1.13
+openai==2.32.0
+openapi-pydantic==0.5.1
+openenv-core==0.2.3
+opentelemetry-api==1.41.1
+orjson==3.11.8
 packaging==26.2
 pandas==3.0.2
+pathable==0.5.0
+pillow==12.2.0
 platformdirs==4.9.6
 propcache==0.4.1
 protobuf==7.34.1
 psutil==7.2.2
+py-key-value-aio==0.4.4
 pyarrow==24.0.0
+pycparser==3.0
 pydantic==2.13.3
+pydantic-settings==2.14.0
 pydantic_core==2.46.3
+pydub==0.25.1
 Pygments==2.20.0
+PyJWT==2.12.1
+pyperclip==1.11.0
 python-dateutil==2.9.0.post0
 python-dotenv==1.2.2
+python-multipart==0.0.26
+pytz==2026.1.post1
 PyYAML==6.0.3
+referencing==0.37.0
 regex==2026.4.4
 requests==2.33.1
 rich==15.0.0
+rich-rst==1.3.2
+rpds-py==0.30.0
+safehttpx==0.1.7
 safetensors==0.7.0
+semantic-version==2.10.0
 sentry-sdk==2.58.0
 setuptools==81.0.0
 shellingham==1.5.4
 six==1.17.0
 smmap==5.0.3
+sniffio==1.3.1
+sse-starlette==3.3.4
 starlette==1.0.0
 sympy==1.14.0
 tokenizers==0.22.2
+tomli==2.4.1
+tomli_w==1.2.0
+tomlkit==0.14.0
 torch==2.11.0
 tqdm==4.67.3
 transformers==5.6.2
@@ -66,9 +126,12 @@ trl==1.2.0
 typer==0.24.2
 typing-inspection==0.4.2
 typing_extensions==4.15.0
+uncalled-for==0.3.1
 urllib3==2.6.3
 uvicorn==0.46.0
 wandb==0.26.1
+watchfiles==1.1.1
 websockets==16.0
 xxhash==3.6.0
 yarl==1.23.0
+zipp==3.23.1