Update environment.py

environment.py

@@ -152,7 +152,9 @@ class EmailTriageEnv:
 
         base_result = self._grade_current_step(validated_action)
         base_score = base_result.score
-
+        previous_base_score = self._base_score_history[-1] if self._base_score_history else None
+        progress_signal = self._compute_progress_signal(base_score, previous_base_score)
+        
         truth_for_step = (
             self._ground_truth[min(self._current_index, len(self._ground_truth) - 1)]
             if self._ground_truth
@@ -166,8 +168,9 @@ class EmailTriageEnv:
 
         penalties = self._compute_penalties(validated_action)
         trajectory_bonus = self._compute_trajectory_bonus()
+        step_cost = self._compute_step_cost()
         final_reward = self._clip_reward(
-            base_score - (self._current_step * 0.01) + trajectory_bonus - penalties
+            base_score + progress_signal + trajectory_bonus - penalties - step_cost
         )
 
         self._reward_history.append(final_reward)
@@ -193,6 +196,8 @@ class EmailTriageEnv:
             "emails_processed": min(self._current_index, len(self._emails)),
             "emails_remaining": max(len(self._emails) - self._current_index, 0),
             "base_score": round(base_score, 4),
+            "progress_signal": round(progress_signal, 4),
+            "step_cost": round(step_cost, 4),
             "penalties": round(penalties, 4),
             "trajectory_bonus": round(trajectory_bonus, 4),
             "grading_feedback": base_result.feedback,
@@ -382,6 +387,39 @@ class EmailTriageEnv:
 
         return penalty_total
 
+        def _compute_progress_signal(
+        self,
+        base_score: float,
+        previous_base_score: float | None,
+    ) -> float:
+        """Compute dense partial-progress reward independent of final completion.
+
+        Args:
+            base_score: Current-step base grade in [0.0, 1.0].
+            previous_base_score: Previous step base grade when available.
+
+        Returns:
+            Small positive/negative signal reflecting progress and quality trend.
+        """
+        total_emails = max(len(self._emails), 1)
+        progress_ratio = min(1.0, (self._current_index + 1) / total_emails)
+
+        completion_signal = 0.05 * progress_ratio
+        quality_signal = 0.05 * self._clip_reward(base_score)
+
+        trend_signal = 0.0
+        if previous_base_score is not None:
+            delta = base_score - previous_base_score
+            trend_signal = max(-0.02, min(0.03, delta * 0.1))
+
+        return completion_signal + quality_signal + trend_signal
+
+    def _compute_step_cost(self) -> float:
+        """Return a gentle efficiency cost that grows with episode length."""
+        normalized_step = self._current_step / max(self._max_steps, 1)
+        return 0.005 + (0.01 * normalized_step)
+
+
     def _compute_trajectory_bonus(self) -> float:
         """Return trajectory bonus when episode completion quality is high.