Spaces:
Sleeping
Sleeping
| """ | |
| Reward Function for MediGuard-AI | |
| ================================ | |
| Maps (action × patient_condition × activity) to a reward signal with: | |
| - Base reward table | |
| - Activity context multipliers (the novel mechanic) | |
| - Alarm fatigue modifier | |
| - Personalization bonus | |
| """ | |
| from enum import Enum | |
| class Action(Enum): | |
| IGNORE = 0 | |
| VERIFY = 1 | |
| ALERT = 2 | |
| class PatientCondition(Enum): | |
| STABLE = "stable" | |
| BORDERLINE = "borderline" | |
| EMERGENCY = "emergency" | |
| DRUG_MASKED = "drug_masked" | |
| # Base reward table: REWARD_TABLE[action][condition] | |
| REWARD_TABLE = { | |
| Action.ALERT: { | |
| PatientCondition.EMERGENCY: +1.0, | |
| PatientCondition.BORDERLINE: +0.2, | |
| PatientCondition.STABLE: -0.5, | |
| PatientCondition.DRUG_MASKED: +1.0, | |
| }, | |
| Action.VERIFY: { | |
| PatientCondition.EMERGENCY: +0.3, | |
| PatientCondition.BORDERLINE: +0.7, | |
| PatientCondition.STABLE: -0.1, | |
| PatientCondition.DRUG_MASKED: +0.3, | |
| }, | |
| Action.IGNORE: { | |
| PatientCondition.EMERGENCY: -1.0, | |
| PatientCondition.BORDERLINE: -0.2, | |
| PatientCondition.STABLE: +0.2, | |
| PatientCondition.DRUG_MASKED: -1.0, | |
| }, | |
| } | |
| # Activity context multipliers — the key insight: | |
| # Same vital sign means different things depending on what the patient is doing. | |
| # HR 130 while walking → expected (low multiplier, discount the anomaly) | |
| # HR 130 while lying still → emergency (high multiplier, amplify the anomaly) | |
| ACTIVITY_CONTEXT = { | |
| 0: 1.00, # resting (lying in bed) — baseline, no discount | |
| 1: 0.40, # eating — slight HR/BP increase expected | |
| 2: 0.50, # walking/ambulating — elevated vitals expected | |
| 3: 1.25, # distressed — amplify concern | |
| 4: 1.60, # falling — immediate concern | |
| } | |
| # Alarm fatigue settings | |
| FATIGUE_WINDOW = 30 # look at last 30 steps | |
| FATIGUE_THRESHOLD = 5 # more than 5 alerts in the window | |
| FATIGUE_MULTIPLIER = 0.6 # reduce reward to 60% | |
| # Personalization settings | |
| PERSONALIZATION_STEP = 200 # kicks in after this many steps | |
| PERSONALIZATION_BONUS = 0.2 # bonus for correctly ignoring known-normal | |
| class RewardFunction: | |
| """ | |
| Stateful reward calculator. | |
| Tracks action history for alarm fatigue and personalization bonuses. | |
| """ | |
| def __init__(self): | |
| self.action_history = [] | |
| self.condition_history = [] | |
| self.activity_history = [] | |
| self.step_count = 0 | |
| def reset(self): | |
| """Clear all history for a new episode.""" | |
| self.action_history = [] | |
| self.condition_history = [] | |
| self.activity_history = [] | |
| self.step_count = 0 | |
| def compute(self, action: Action, condition: PatientCondition, | |
| activity: int = 0) -> float: | |
| """ | |
| Compute the reward for a single (action, condition, activity) tuple. | |
| Parameters | |
| ---------- | |
| action : Action | |
| The agent's action (IGNORE, VERIFY, ALERT). | |
| condition : PatientCondition | |
| The patient's current condition. | |
| activity : int | |
| The patient's current activity code (0-4). | |
| Returns | |
| ------- | |
| float | |
| The reward value (can be negative). | |
| """ | |
| self.step_count += 1 | |
| self.action_history.append(action) | |
| self.condition_history.append(condition) | |
| self.activity_history.append(activity) | |
| # 1. Base reward from the table | |
| base_reward = REWARD_TABLE[action][condition] | |
| # 2. Activity context multiplier | |
| # Only apply to penalty situations — don't discount correct actions | |
| ctx = ACTIVITY_CONTEXT.get(activity, 1.0) | |
| if base_reward < 0: | |
| # Penalties are reduced during expected-high-vitals activities | |
| # e.g., alerting during walking gets less penalty (ctx=0.5) | |
| base_reward *= ctx | |
| elif condition in (PatientCondition.EMERGENCY, PatientCondition.DRUG_MASKED): | |
| # Correct emergency responses amplified during dangerous activities | |
| base_reward *= ctx | |
| # 3. Alarm fatigue modifier | |
| fatigue_modifier = 1.0 | |
| if len(self.action_history) >= FATIGUE_WINDOW: | |
| recent_alerts = sum( | |
| 1 for a in self.action_history[-FATIGUE_WINDOW:] | |
| if a == Action.ALERT | |
| ) | |
| if recent_alerts > FATIGUE_THRESHOLD: | |
| fatigue_modifier = FATIGUE_MULTIPLIER | |
| # 4. Personalization bonus | |
| personalization = 0.0 | |
| if (self.step_count > PERSONALIZATION_STEP | |
| and action == Action.IGNORE | |
| and condition == PatientCondition.STABLE): | |
| personalization = PERSONALIZATION_BONUS | |
| # Final reward | |
| reward = base_reward * fatigue_modifier + personalization | |
| return reward | |
| def get_stats(self) -> dict: | |
| """Return episode statistics for graders.""" | |
| total_alerts = sum(1 for a in self.action_history if a == Action.ALERT) | |
| total_verifies = sum(1 for a in self.action_history if a == Action.VERIFY) | |
| total_ignores = sum(1 for a in self.action_history if a == Action.IGNORE) | |
| return { | |
| "total_steps": self.step_count, | |
| "total_alerts": total_alerts, | |
| "total_verifies": total_verifies, | |
| "total_ignores": total_ignores, | |
| "action_history": list(self.action_history), | |
| "condition_history": list(self.condition_history), | |
| "activity_history": list(self.activity_history), | |
| } | |