Clamp step scores with SAFE_SCORE

Author: Gowwwthami

app.py

@@ -14,6 +14,7 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, Field
 from openai import OpenAI
+from src.score_utils import SAFE_SCORE
 
 
 # ============================================================================
@@ -144,7 +145,7 @@ def step(self, action: str) -> dict:
         if self.done:
             return {
                 "state": "environment already done",
-                "reward": 0.0,
+                "reward": SAFE_SCORE(0.0),
                 "done": True
             }
 
@@ -157,7 +158,7 @@ def step(self, action: str) -> dict:
         else:
             return {
                 "state": f"invalid action: {action}",
-                "reward": 0.0,
+                "reward": SAFE_SCORE(0.0),
                 "done": self.done
             }
 
@@ -170,7 +171,7 @@ def _classify_email(self) -> dict:
 
         return {
             "state": f"classified as: {self.category}",
-            "reward": 0.33,
+            "reward": SAFE_SCORE(0.33),
             "done": False,
             "info": {
                 "action": "classify_email",
@@ -197,7 +198,7 @@ def _extract_entities(self) -> dict:
 
         return {
             "state": f"extracted: {json.dumps(self.extracted_data)}",
-            "reward": 0.33,
+            "reward": SAFE_SCORE(0.33),
             "done": False,
             "info": {
                 "action": "extract_entities",
@@ -229,7 +230,7 @@ def _generate_reply(self) -> dict:
 
         return {
             "state": f"reply generated: {self.response[:100]}...",
-            "reward": 0.34,
+            "reward": SAFE_SCORE(0.34),
             "done": True,
             "info": {
                 "action": "generate_reply",

server/app.py

@@ -21,7 +21,7 @@
 
 from src.env import EmailTriageEnv, OpenEnvEmailTriageEnv
 from src.models import Action
-    from src.score_utils import SAFE_SCORE
+from src.score_utils import SAFE_SCORE
 
 
 ALLOWED_CATEGORIES = {"billing", "technical", "sales", "account", "complaint", "shipping", "other"}
@@ -77,7 +77,7 @@
 def _new_scoreboard() -> Dict[str, Dict[str, Any]]:
     return {
         task_id: {
-            "score": 0.0,
+            "score": SAFE_SCORE(0.01),
             "steps": 0,
             "cumulative_reward": 0.0,
             "done": False,

src/env.py

@@ -63,6 +63,7 @@ def step(self, action: Action | Dict[str, Any]) -> Tuple[Observation, float, boo
 
         current_email = self.dataset[self.index]
         reward_obj = compute_step_reward(action=action, truth=current_email, task=self.task)
+        reward = SAFE_SCORE(reward_obj.total)
         self.grader.update(action=action, truth=current_email)
 
         # Update tracking statistics
@@ -77,8 +78,8 @@ def step(self, action: Action | Dict[str, Any]) -> Tuple[Observation, float, boo
         if action.priority and action.priority.value == "urgent":
             self.urgent_handled += 1
 
-        self.last_reward = reward_obj.total
-        self.cumulative_reward += reward_obj.total
+        self.last_reward = reward
+        self.cumulative_reward += reward
 
         self.index += 1
         self.done = self.index >= len(self.dataset)
@@ -106,7 +107,7 @@ def step(self, action: Action | Dict[str, Any]) -> Tuple[Observation, float, boo
                 "urgent_handled": self.urgent_handled,
             },
         }
-        return next_observation, reward_obj.total, self.done, info
+        return next_observation, reward, self.done, info
 
     def final_score(self) -> float:
         return SAFE_SCORE(self.grader.score())
@@ -142,7 +143,7 @@ def step(self, action: Dict[str, Any]) -> Tuple[Dict[str, Any], float, bool, Dic
         predicted_action = str(action.get("action", "")).strip()
         expected_action = self._expected_action_for_email(current_email)
         correct = predicted_action == expected_action
-        reward = 1.0 if correct else -1.0
+        reward = SAFE_SCORE(1.0 if correct else -1.0)
 
         self.current_index += 1
         self.done = self.current_index >= len(self.emails)

src/graders.py

@@ -6,7 +6,8 @@
 
 
 def safe_score(correct: int, total: int) -> float:
-    return SAFE_SCORE(safe_ratio_score(correct=correct, total=total))
+    score = safe_ratio_score(correct=correct, total=total)
+    return SAFE_SCORE(score)
 
 
 class DeterministicTriageGrader:
@@ -37,4 +38,5 @@ def update(self, action: Action, truth: EmailRecord) -> None:
                 self.correct += 1
 
     def score(self) -> float:
-        return SAFE_SCORE(safe_ratio_score(correct=self.correct, total=self.total))
+        score = safe_ratio_score(correct=self.correct, total=self.total)
+        return SAFE_SCORE(score)

src/rewards.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from src.models import Action, EmailRecord, Reward, TriageAction
+from src.score_utils import SAFE_SCORE
 from src.tasks import TaskConfig
 
 
@@ -39,6 +40,7 @@ def compute_step_reward(action: Action, truth: EmailRecord, task: TaskConfig) ->
         penalties["unnecessary_escalation"] = UNNECESSARY_ESCALATION_PENALTY
 
     total = category_component + priority_component + action_component + reply_component + sum(penalties.values())
+    total = SAFE_SCORE(total)
 
     return Reward(
         total=total,