Add /grade endpoint, use real grader scores in evaluator

baseline/evaluator.py

@@ -20,7 +20,9 @@ from core.trajectory import EpisodicMemory
 from core.reward import normalise_score
 from core.policy_update import compute_advantage, update_memory
 
-N_ROLLOUTS = 3
+import requests as http_requests
+
+N_ROLLOUTS = 4
 
 
 # ── Prompt Builder With Memory ────────────────────────────────────────────────
@@ -46,11 +48,11 @@ def build_prompt_with_memory(obs: CityObservation, memory: EpisodicMemory) -> st
 
 def run_rollout(
     env:       Any,
-    task_name: str,         
+    task_name: str,
     client:    OpenAI,
     memory:    EpisodicMemory,
     verbose:   bool = True,
-    ) -> Tuple[float, int, List[dict]]: 
+) -> Tuple[float, int, List[dict]]:
     """Run one complete episode using memory-augmented prompts."""
     result = env.reset(task_name=task_name)
     obs          = result.observation
@@ -64,7 +66,15 @@ def run_rollout(
         response = call_llm(prompt, client)
         action   = parse_action(response, len(obs.districts))
 
-        result       = env.step(action)
+        try:
+            result   = env.step(action)
+        except Exception as e:
+            if "close frame" in str(e).lower() or "websocket" in str(e).lower():
+                if verbose:
+                    print(f"      ⚠ WebSocket dropped at step {step+1}, ending rollout early")
+                break
+            raise
+
         next_obs     = result.observation
         reward       = result.reward or 0.0
         done         = result.done
@@ -97,6 +107,7 @@ def run_task_grpo(
     env:       Any,
     task_name: str,
     client:    OpenAI,
+    base_url:  str,
     verbose:   bool = True,
 ) -> float:
     """GRPO-style simulated learning loop for one task."""
@@ -113,11 +124,27 @@ def run_task_grpo(
             print(f"\n    Rollout {i+1}/{N_ROLLOUTS} [{label}]")
 
         total_reward, steps, trajectory = run_rollout(env, task_name, client, memory, verbose)
-        score                           = normalise_score(total_reward, steps)
-        rollouts.append((total_reward, steps, score))
 
-        if verbose:
-            print(f"    → Reward: {total_reward:+.4f} | Score: {score:.4f}")
+        # Use real grader score from server
+        num_districts = {"easy": 2, "medium": 4, "hard": 6}.get(task_name, 2)
+        try:
+            grade_resp = http_requests.get(
+                base_url.rstrip('/') + '/grade', timeout=10
+            )
+            if grade_resp.status_code == 200:
+                data  = grade_resp.json()
+                score = data["final_score"]
+                if verbose:
+                    print(
+                        f"    → Grader: containment={data['containment_score']:.3f} "
+                        f"hospital={data['hospital_score']:.3f} "
+                        f"efficiency={data['efficiency_score']:.3f} "
+                        f"speed={data['speed_score']:.3f}"
+                    )
+            else:
+                score = normalise_score(total_reward, steps, num_districts)
+        except Exception:
+            score = normalise_score(total_reward, steps, num_districts)
 
         # GRPO advantage computation
         completed_rewards = [r[0] for r in rollouts]
@@ -164,7 +191,7 @@ def run_evaluation(
     with CascadeContainmentEnv(base_url=base_url).sync() as env:
         for task_name in ["easy", "medium", "hard"]:
             try:
-                score             = run_task_grpo(env, task_name, client, verbose)
+                score = run_task_grpo(env, task_name, client, base_url, verbose)
                 scores[task_name] = score
                 if verbose:
                     print(f"\n  ✓ {task_name.upper()} final score: {score:.4f}")

server/app.py

@@ -10,9 +10,10 @@ import os
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 
 from openenv.core.env_server import create_app
+from fastapi.responses import JSONResponse
 from server.environment import EpidemicContainmentEnv
 from models import ContainmentAction, CityObservation
-
+import server.environment as env_module
 
 app = create_app(
     EpidemicContainmentEnv,
@@ -20,6 +21,19 @@ app = create_app(
     CityObservation,
 )
 
+@app.get("/grade")
+async def grade_last_episode():
+    """
+    Returns the deterministic grader score for the most recently completed episode.
+    Computed automatically when an episode ends via the WebSocket session.
+    """
+    if not env_module._last_grade:
+        return JSONResponse(
+            {"error": "No completed episode yet — run a full episode first"},
+            status_code=400
+        )
+    return JSONResponse(env_module._last_grade)
+
 from fastapi.responses import HTMLResponse
 
 @app.get("/", response_class=HTMLResponse)

server/environment.py

@@ -52,6 +52,8 @@ from server.utils import (
 )
 from server.tasks.registry import get_task
 
+from server.grader import grade_trajectory
+_last_grade: dict = {}
 
 class EpidemicContainmentEnv(Environment):
     """
@@ -159,6 +161,21 @@ class EpidemicContainmentEnv(Environment):
         # ── 10. Check terminal conditions ─────────────────────────────────────
         done, terminal_message = self._check_terminal()
 
+        if done and self._trajectory:
+            import server.environment as _self_module
+            result = grade_trajectory(self._trajectory, self._task_name)
+            _self_module._last_grade = {
+                "final_score":         result.final_score,
+                "containment_score":   result.containment_score,
+                "hospital_score":      result.hospital_score,
+                "efficiency_score":    result.efficiency_score,
+                "speed_score":         result.speed_score,
+                "hospital_breached":   result.hospital_breached,
+                "districts_contained": result.districts_contained,
+                "total_steps":         result.total_steps,
+                "task_name":           self._task_name,
+            }
+
         # ── 11. Build and return observation ──────────────────────────────────
         final_message = terminal_message if terminal_message else message