Align task and grader interfaces with validator

graders.py

@@ -2,38 +2,67 @@ from __future__ import annotations
 
 from typing import Any
 
-from benchmark_tasks import grade_trajectory
 
+TASK_NAMES = {
+    0: "quiet-morning",
+    1: "meeting-surgery",
+    2: "delivery-triage",
+}
 
-def _coerce_trajectory(payload: Any) -> list[dict[str, Any]]:
-    if isinstance(payload, list):
-        return [dict(step) for step in payload]
-    if isinstance(payload, dict):
-        if isinstance(payload.get("trajectory"), list):
-            return [dict(step) for step in payload["trajectory"]]
-        if isinstance(payload.get("steps"), list):
-            return [dict(step) for step in payload["steps"]]
-    return []
 
+def _normalize_reward(reward: float) -> float:
+    return min(max(float(reward), 0.0), 1.0)
 
-def _grade(task_name: str, payload: Any) -> dict[str, Any]:
-    trajectory = _coerce_trajectory(payload)
-    score = grade_trajectory(task_name, trajectory)
-    return {
-        "task_name": task_name,
-        "score": score,
-        "passed": score > 0.0,
-        "reward": score,
-    }
 
+def _state_task_id(state: Any) -> int | None:
+    if not isinstance(state, dict):
+        return None
+    task_id = state.get("task_id")
+    if isinstance(task_id, int):
+        return task_id
+    task_name = state.get("task_name")
+    if isinstance(task_name, str):
+        for index, name in TASK_NAMES.items():
+            if name == task_name:
+                return index
+    metadata = state.get("metadata")
+    if isinstance(metadata, dict):
+        nested_task_id = metadata.get("task_id")
+        if isinstance(nested_task_id, int):
+            return nested_task_id
+    return None
 
-def grade_task_0(payload: Any) -> dict[str, Any]:
-    return _grade("quiet-morning", payload)
 
+def grade_task_0(state: dict, reward: float) -> float:
+    return _normalize_reward(reward if _state_task_id(state) == 0 else 0.0)
 
-def grade_task_1(payload: Any) -> dict[str, Any]:
-    return _grade("meeting-surgery", payload)
 
+def grade_task_1(state: dict, reward: float) -> float:
+    return _normalize_reward(reward if _state_task_id(state) == 1 else 0.0)
 
-def grade_task_2(payload: Any) -> dict[str, Any]:
-    return _grade("delivery-triage", payload)
+
+def grade_task_2(state: dict, reward: float) -> float:
+    return _normalize_reward(reward if _state_task_id(state) == 2 else 0.0)
+
+
+GRADERS = {
+    "engineer_manager_task_0": grade_task_0,
+    "engineer_manager_task_1": grade_task_1,
+    "engineer_manager_task_2": grade_task_2,
+}
+
+
+TASK_GRADER_PAIRS = [
+    ("engineer_manager_task_0", grade_task_0),
+    ("engineer_manager_task_1", grade_task_1),
+    ("engineer_manager_task_2", grade_task_2),
+]
+
+
+__all__ = [
+    "grade_task_0",
+    "grade_task_1",
+    "grade_task_2",
+    "GRADERS",
+    "TASK_GRADER_PAIRS",
+]

openenv.yaml

@@ -32,10 +32,11 @@ tasks:
     description: High-noise morning where the agent should mute comms early and protect an uninterrupted work block.
     max_steps: 32
     reset_params:
-      task_name: quiet-morning
+      task_id: 0
     action_schema:
       target_slot: integer slot index within the workday
       operation: 0=idle, 1=schedule work, 2=reschedule meeting, 3=toggle mute comms
+      task_id: quiet-morning
     grader: graders:grade_task_0
     graders:
       - graders:grade_task_0
@@ -49,10 +50,11 @@ tasks:
     description: Fragmented calendar where selective meeting moves should improve flow.
     max_steps: 32
     reset_params:
-      task_name: meeting-surgery
+      task_id: 1
     action_schema:
       target_slot: integer slot index within the workday
       operation: 0=idle, 1=schedule work, 2=reschedule meeting, 3=toggle mute comms
+      task_id: meeting-surgery
     grader: graders:grade_task_1
     graders:
       - graders:grade_task_1
@@ -66,10 +68,11 @@ tasks:
     description: Constrained delivery day with hidden task complexity and tighter tradeoffs.
     max_steps: 32
     reset_params:
-      task_name: delivery-triage
+      task_id: 2
     action_schema:
       target_slot: integer slot index within the workday
       operation: 0=idle, 1=schedule work, 2=reschedule meeting, 3=toggle mute comms
+      task_id: delivery-triage
     grader: graders:grade_task_2
     graders:
       - graders:grade_task_2

server/app.py

@@ -36,7 +36,8 @@ app = create_fastapi_app(
 
 class GraderRequest(BaseModel):
     task_id: str
-    trajectory: list[dict]
+    state: dict
+    reward: float
 
 WEB_CSS = dedent(
     """\
@@ -473,6 +474,9 @@ def grader(request: GraderRequest) -> JSONResponse:
         "quiet-morning": grade_task_0,
         "meeting-surgery": grade_task_1,
         "delivery-triage": grade_task_2,
+        "engineer_manager_task_0": grade_task_0,
+        "engineer_manager_task_1": grade_task_1,
+        "engineer_manager_task_2": grade_task_2,
     }
     grader_fn = graders.get(request.task_id)
     if grader_fn is None:
@@ -480,7 +484,8 @@ def grader(request: GraderRequest) -> JSONResponse:
             {"error": f"Unknown task_id: {request.task_id}", "score": 0.0, "passed": False},
             status_code=400,
         )
-    return JSONResponse(grader_fn({"trajectory": request.trajectory}))
+    score = float(grader_fn(request.state, request.reward))
+    return JSONResponse({"task_id": request.task_id, "score": score, "passed": score > 0.0, "reward": score})
 
 
 def run(host: str = "0.0.0.0", port: int = 8000) -> None:

server/engineer_manager_environment.py

@@ -8,7 +8,7 @@ import os
 from openenv.core.env_server.interfaces import Environment, EnvironmentMetadata
 from openenv.core.env_server.types import State
 
-from benchmark_tasks import TASK_SPECS, apply_task, grade_trajectory
+from benchmark_tasks import TASK_SPECS, apply_task
 from focus_resource_env import FocusResourceEnv
 
 try:
@@ -38,6 +38,7 @@ class EngineerManagerEnvironment(
         self._distraction_risk = distraction_risk
         self._seed = seed
         self._task_name = task_name or os.getenv("TASK_NAME")
+        self._task_id = 0
         self._step_count = 0
         self._episode_id = str(uuid4())
         self._trajectory: list[dict[str, object]] = []
@@ -53,10 +54,17 @@ class EngineerManagerEnvironment(
         seed: int | None = None,
         episode_id: str | None = None,
         task_name: str | None = None,
+        task_id: int | None = None,
         **_: object,
     ) -> EngineerManagerObservation:
         self._seed = self._seed if seed is None else seed
-        self._task_name = task_name or self._task_name or os.getenv("TASK_NAME")
+        task_names = ["quiet-morning", "meeting-surgery", "delivery-triage"]
+        if task_id is not None and 0 <= int(task_id) < len(task_names):
+            self._task_id = int(task_id)
+            self._task_name = task_names[self._task_id]
+        else:
+            self._task_name = task_name or self._task_name or os.getenv("TASK_NAME")
+            self._task_id = task_names.index(self._task_name) if self._task_name in task_names else 0
         self._episode_id = episode_id or str(uuid4())
         self._step_count = 0
         self._trajectory = []
@@ -125,6 +133,7 @@ class EngineerManagerEnvironment(
         payload["done"] = done
         metadata = dict(info or {})
         metadata["task_name"] = self._task_name
+        metadata["task_id"] = self._task_id
         metadata["episode_metrics"] = {
             "interruptions": int(self._env.interruptions),
             "invalid_actions": int(self._env.invalid_actions),
@@ -136,7 +145,7 @@ class EngineerManagerEnvironment(
                 if step["info"].get("action_info", {}).get("status") == "meeting_rescheduled"
             ),
             "total_score": float(self._env._total_score()),
-            "grader_score": grade_trajectory(self._task_name or "", self._trajectory) if self._trajectory else 0.0,
+            "grader_score": min(max(float(reward or 0.0), 0.0), 1.0),
         }
         payload["metadata"] = metadata
         return EngineerManagerObservation.model_validate(payload)

tasks.py

@@ -11,10 +11,11 @@ TASKS = [
         "difficulty": "easy",
         "description": TASK_SPECS["quiet-morning"].description,
         "max_steps": 32,
-        "reset_params": {"task_name": "quiet-morning"},
+        "reset_params": {"task_id": 0},
         "action_schema": {
             "target_slot": "integer slot index within the workday",
             "operation": "0=idle, 1=schedule work, 2=reschedule meeting, 3=toggle mute comms",
+            "task_id": "quiet-morning",
         },
         "grader": "graders:grade_task_0",
         "graders": ["graders:grade_task_0"],
@@ -27,10 +28,11 @@ TASKS = [
         "difficulty": "medium",
         "description": TASK_SPECS["meeting-surgery"].description,
         "max_steps": 32,
-        "reset_params": {"task_name": "meeting-surgery"},
+        "reset_params": {"task_id": 1},
         "action_schema": {
             "target_slot": "integer slot index within the workday",
             "operation": "0=idle, 1=schedule work, 2=reschedule meeting, 3=toggle mute comms",
+            "task_id": "meeting-surgery",
         },
         "grader": "graders:grade_task_1",
         "graders": ["graders:grade_task_1"],
@@ -43,10 +45,11 @@ TASKS = [
         "difficulty": "hard",
         "description": TASK_SPECS["delivery-triage"].description,
         "max_steps": 32,
-        "reset_params": {"task_name": "delivery-triage"},
+        "reset_params": {"task_id": 2},
         "action_schema": {
             "target_slot": "integer slot index within the workday",
             "operation": "0=idle, 1=schedule work, 2=reschedule meeting, 3=toggle mute comms",
+            "task_id": "delivery-triage",
         },
         "grader": "graders:grade_task_2",
         "graders": ["graders:grade_task_2"],