Unify OpenEnv and benchmark episode scoring

src/benchmark.py

@@ -7,8 +7,8 @@ import random
 from typing import Any
 
 from src.models import Action, DispatchAction
+from src.grading import grade_episode
 from src.openenv_environment import OpenEnvEnvironment
-from src.rewards import TaskGrader
 from src.tasks.registry import TaskRegistry
 
 
@@ -70,9 +70,8 @@ async def _run_episode_async(task_id: str, seed: int) -> tuple[float, list[float
             metadata={},
         )
 
-    # Score episodes the same way as the OpenEnv evaluation path:
-    # a normalized aggregate of per-step rewards.
-    final_score = TaskGrader().grade_episode(rewards, task_id=task_id)
+    # Score episodes the same way as the OpenEnv evaluation path.
+    final_score = grade_episode(task_id=task_id, state=final_state, rewards=rewards)
     return final_score, rewards
 
 

src/grading.py

@@ -0,0 +1,49 @@
+"""Episode grading utilities.
+
+This module centralizes "final score" computation so benchmark runs and
+OpenEnv runs report the same episode score.
+"""
+
+from __future__ import annotations
+
+from src.models import State
+
+
+def grade_episode(task_id: str, state: State | None, rewards: list[float]) -> float:
+    """Compute a final episode score in [0.0, 1.0].
+
+    Args:
+        task_id: Task identifier.
+        state: Final (or current) state.
+        rewards: Per-step rewards.
+
+    Returns:
+        Normalized score in [0.0, 1.0].
+    """
+
+    if not rewards:
+        return 0.0
+
+    # Lazy imports avoid circular dependencies (task graders import src.rewards).
+    if task_id == "single_incident":
+        from src.tasks.single_incident import SingleIncidentGrader
+
+        return float(SingleIncidentGrader().grade(state, rewards) if state is not None else 0.0)
+
+    if task_id == "multi_incident":
+        from src.tasks.multi_incident import MultiIncidentGrader
+
+        return float(MultiIncidentGrader().grade(state, rewards) if state is not None else 0.0)
+
+    if task_id == "mass_casualty":
+        from src.tasks.mass_casualty import MassCasualtyGrader
+
+        return float(MassCasualtyGrader().grade(state, rewards) if state is not None else 0.0)
+
+    if task_id == "shift_surge":
+        from src.tasks.shift_surge import ShiftSurgeGrader
+
+        return float(ShiftSurgeGrader().grade(state, rewards) if state is not None else 0.0)
+
+    # Fallback: mean of rewards (legacy behavior).
+    return float(sum(rewards) / max(len(rewards), 1))

src/openenv_environment.py

@@ -3,6 +3,7 @@
 import uuid
 
 from src.city_schema import CitySchemaLoader
+from src.grading import grade_episode
 from src.models import Action, Observation, State
 from src.state_machine import DispatchStateMachine
 
@@ -20,6 +21,8 @@ class OpenEnvEnvironment:
         episode_id = str(uuid.uuid4())
         self._state = self._machine.reset(task_id=self.task_id, episode_id=episode_id)
         self._state.metadata["cumulative_reward"] = 0.0
+        self._state.metadata["episode_rewards"] = []
+        self._state.metadata["episode_score"] = 0.0
         self._last_observation = Observation(
             result="dispatch center online",
             score=0.0,
@@ -40,12 +43,29 @@ class OpenEnvEnvironment:
             raise RuntimeError("Environment not initialized. Call reset() first.")
         state, obs = self._machine.step(self._state, action)
         self._state = state
-        self._last_observation = obs
+
+        # `DispatchStateMachine.step()` sets `obs.score` to the per-step reward.
+        # OpenEnv consumers often interpret `observation.score` as an episode score,
+        # so we keep the per-step reward in `reward` and publish the episode score
+        # into `observation.score`.
+        step_reward = float(obs.score)
+
+        rewards: list[float] = list(self._state.metadata.get("episode_rewards", []))
+        rewards.append(step_reward)
+        self._state.metadata["episode_rewards"] = rewards
+
         cumulative = float(self._state.metadata.get("cumulative_reward", 0.0))
-        cumulative += float(obs.score)
-        self._state.metadata["cumulative_reward"] = cumulative
+        self._state.metadata["cumulative_reward"] = cumulative + step_reward
+
+        # Episode score is derived from the same grading logic as benchmark runs.
+        episode_score = grade_episode(task_id=self.task_id, state=self._state, rewards=rewards)
+        episode_score = max(0.0, min(1.0, float(episode_score)))
+        self._state.metadata["episode_score"] = episode_score
+
         done = self._machine.is_terminal(state)
-        return obs, obs.score, done
+        obs = obs.model_copy(update={"score": episode_score})
+        self._last_observation = obs
+        return obs, step_reward, done
 
     def state(self) -> State:
         if self._state is None:

tests/test_benchmark_integration.py

@@ -2,7 +2,11 @@
 
 from __future__ import annotations
 
+import asyncio
+
+from src.models import Action, DispatchAction
 from src.benchmark import list_tasks, run_all, run_task
+from src.openenv_environment import OpenEnvEnvironment
 
 
 def test_list_tasks_has_four() -> None:
@@ -16,13 +20,34 @@ def test_run_task_score_in_range() -> None:
     result = run_task("single_incident", seed=42)
     assert 0.0 <= result["score"] <= 1.0
     assert result["task_id"] == "single_incident"
-    # Benchmark scoring must match the OpenEnv evaluation path: mean step reward.
-    rewards = result["rewards"]
-    if rewards:
-        expected = sum(rewards) / len(rewards)
-    else:
-        expected = 0.0
-    assert abs(result["score"] - expected) < 1e-9
+
+
+def test_benchmark_and_openenv_use_same_episode_grader(monkeypatch) -> None:
+    from src.tasks.single_incident import SingleIncidentGrader
+
+    expected_score = 0.777
+    monkeypatch.setattr(SingleIncidentGrader, "grade", lambda self, state, rewards: expected_score)
+
+    # Benchmark path.
+    result = run_task("single_incident", seed=42)
+    assert abs(result["score"] - expected_score) < 1e-9
+
+    # OpenEnv path.
+    env = OpenEnvEnvironment(task_id="single_incident", seed=42)
+    asyncio.run(env.reset())
+    obs, reward, done = asyncio.run(
+        env.step(
+            Action(
+                action_type=DispatchAction.DISPATCH,
+                unit_id="MED-1",
+                incident_id="INC-001",
+            )
+        )
+    )
+    assert isinstance(reward, float)
+    assert isinstance(done, bool)
+    assert abs(float(obs.score) - expected_score) < 1e-9
+    env.close()
 
 
 def test_run_all_scores_in_range() -> None: