Add episode-level reward helper for RL training

rl/rewards.py

@@ -0,0 +1,127 @@
+"""Shared reward helper utilities for RL agents.
+
+The helper operates on episode-level statistics so that reward shaping
+reflects system-wide outcomes (disposal rate, gap compliance, urgent
+case latency, and fairness across cases).
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+from typing import Dict, Iterable, Optional
+
+import numpy as np
+
+from scheduler.core.case import Case
+
+
+@dataclass
+class EpisodeRewardHelper:
+    """Aggregates episode metrics and computes shaped rewards."""
+
+    total_cases: int
+    target_gap_days: int = 30
+    max_urgent_latency: int = 60
+    disposal_weight: float = 4.0
+    gap_weight: float = 1.5
+    urgent_weight: float = 2.0
+    fairness_weight: float = 1.0
+    _disposed_cases: int = 0
+    _hearing_counts: Dict[str, int] = field(default_factory=lambda: defaultdict(int))
+    _urgent_latencies: list[float] = field(default_factory=list)
+
+    def _base_outcome_reward(self, case: Case, was_scheduled: bool, hearing_outcome: str) -> float:
+        """Preserve the original per-case shaping signals."""
+
+        reward = 0.0
+        if not was_scheduled:
+            return reward
+
+        # Base scheduling reward (small positive for taking action)
+        reward += 0.5
+
+        # Hearing outcome rewards
+        lower_outcome = hearing_outcome.lower()
+        if "disposal" in lower_outcome or "judgment" in lower_outcome or "settlement" in lower_outcome:
+            reward += 10.0  # Major positive for disposal
+        elif "progress" in lower_outcome and "adjourn" not in lower_outcome:
+            reward += 3.0  # Progress without disposal
+        elif "adjourn" in lower_outcome:
+            reward -= 3.0  # Negative for adjournment
+
+        # Urgency bonus
+        if case.is_urgent:
+            reward += 2.0
+
+        # Ripeness penalty
+        if hasattr(case, "ripeness_status") and case.ripeness_status not in ["RIPE", "UNKNOWN"]:
+            reward -= 4.0
+
+        # Long pending bonus (>365 days)
+        if case.age_days and case.age_days > 365:
+            reward += 2.0
+
+        return reward
+
+    def _fairness_score(self) -> float:
+        """Reward higher uniformity in hearing distribution."""
+
+        counts: Iterable[int] = self._hearing_counts.values()
+        if not counts:
+            return 0.0
+
+        counts_array = np.array(list(counts), dtype=float)
+        mean = np.mean(counts_array)
+        if mean == 0:
+            return 0.0
+
+        dispersion = np.std(counts_array) / (mean + 1e-6)
+        # Lower dispersion -> better fairness. Convert to reward in [0, 1].
+        fairness = max(0.0, 1.0 - dispersion)
+        return fairness
+
+    def compute_case_reward(
+        self,
+        case: Case,
+        was_scheduled: bool,
+        hearing_outcome: str,
+        current_date,
+        previous_gap_days: Optional[int] = None,
+    ) -> float:
+        """Compute reward using both local and episode-level signals."""
+
+        reward = self._base_outcome_reward(case, was_scheduled, hearing_outcome)
+
+        if not was_scheduled:
+            return reward
+
+        # Track disposals
+        if "disposal" in hearing_outcome.lower() or getattr(case, "is_disposed", False):
+            self._disposed_cases += 1
+
+        # Track hearing counts for fairness
+        self._hearing_counts[case.case_id] = case.hearing_count or self._hearing_counts[case.case_id] + 1
+
+        # Track urgent latencies
+        if case.is_urgent:
+            self._urgent_latencies.append(case.age_days or 0)
+
+        # Episode-level components
+        disposal_rate = (self._disposed_cases / self.total_cases) if self.total_cases else 0.0
+        reward += self.disposal_weight * disposal_rate
+
+        if previous_gap_days is not None:
+            gap_score = max(0.0, 1.0 - (previous_gap_days / self.target_gap_days))
+            reward += self.gap_weight * gap_score
+
+        if self._urgent_latencies:
+            avg_latency = float(np.mean(self._urgent_latencies))
+            latency_score = max(0.0, 1.0 - (avg_latency / self.max_urgent_latency))
+            reward += self.urgent_weight * latency_score
+
+        fairness = self._fairness_score()
+        reward += self.fairness_weight * fairness
+
+        return reward
+

rl/training.py

@@ -13,7 +13,8 @@ import random
 from scheduler.data.case_generator import CaseGenerator
 from scheduler.simulation.engine import CourtSim, CourtSimConfig
 from scheduler.core.case import Case, CaseStatus
-from .simple_agent import TabularQAgent, CaseState
+from .simple_agent import TabularQAgent
+from .rewards import EpisodeRewardHelper
 
 
 class RLTrainingEnvironment:
@@ -32,6 +33,7 @@ class RLTrainingEnvironment:
         self.horizon_days = horizon_days
         self.current_date = start_date
         self.episode_rewards = []
+        self.reward_helper = EpisodeRewardHelper(total_cases=len(cases))
         
     def reset(self) -> List[Case]:
         """Reset environment for new training episode.
@@ -42,6 +44,7 @@ class RLTrainingEnvironment:
         """
         self.current_date = self.start_date
         self.episode_rewards = []
+        self.reward_helper = EpisodeRewardHelper(total_cases=len(self.cases))
         return self.cases.copy()
     
     def step(self, agent_decisions: Dict[str, int]) -> Tuple[List[Case], Dict[str, float], bool]:
@@ -57,18 +60,23 @@ class RLTrainingEnvironment:
         rewards = {}
         
         # For each case that agent decided to schedule
-        scheduled_cases = [case for case in self.cases 
+        scheduled_cases = [case for case in self.cases
                           if case.case_id in agent_decisions and agent_decisions[case.case_id] == 1]
         
         # Simulate hearing outcomes for scheduled cases
         for case in scheduled_cases:
             if case.is_disposed:
                 continue
-                
+
             # Simulate hearing outcome based on stage transition probabilities
             outcome = self._simulate_hearing_outcome(case)
             was_heard = "heard" in outcome.lower()
-            
+
+            # Track gap relative to previous hearing for reward shaping
+            previous_gap = None
+            if case.last_hearing_date:
+                previous_gap = max(0, (self.current_date - case.last_hearing_date).days)
+
             # Always record the hearing
             case.record_hearing(self.current_date, was_heard=was_heard, outcome=outcome)
             
@@ -83,7 +91,13 @@ class RLTrainingEnvironment:
             # If adjourned, case stays in same stage
             
             # Compute reward for this case
-            rewards[case.case_id] = self._compute_reward(case, outcome)
+            rewards[case.case_id] = self.reward_helper.compute_case_reward(
+                case,
+                was_scheduled=True,
+                hearing_outcome=outcome,
+                current_date=self.current_date,
+                previous_gap_days=previous_gap,
+            )
         
         # Update case ages
         for case in self.cases:
@@ -131,13 +145,7 @@ class RLTrainingEnvironment:
         # Default progression
         return "ARGUMENTS"
     
-    def _compute_reward(self, case: Case, outcome: str) -> float:
-        """Compute reward based on case and outcome."""
-        agent = TabularQAgent()  # Use for reward computation
-        return agent.compute_reward(case, was_scheduled=True, hearing_outcome=outcome)
-
-
-def train_agent(agent: TabularQAgent, episodes: int = 100, 
+def train_agent(agent: TabularQAgent, episodes: int = 100,
                 cases_per_episode: int = 1000, 
                 episode_length: int = 60,
                 verbose: bool = True) -> Dict: