Spaces:
Sleeping
Sleeping
| """ | |
| reward.py | |
| ========= | |
| Deterministic, itemised reward function for the InformationIntegrity environment. | |
| All reward constants are read from REWARD_CONFIG (default: reward_config.yaml) | |
| so platform teams can tune them without touching Python code. | |
| Scoring schema | |
| -------------- | |
| Process rewards (per step): | |
| +verify_claim_reward per unique claim verified | |
| +check_source_reward per unique source checked | |
| +multi_source_bonus when ≥ multi_source_threshold sources consulted | |
| Decision reward (episode end): | |
| +decision_correct decision exactly matches ground truth | |
| +decision_partial decision is one step away (e.g. flag vs remove) | |
| 0 decision is fully wrong | |
| Penalties: | |
| premature_decision deciding before min evidence thresholds are met | |
| redundant_action re-verifying a claim or re-checking a source | |
| step_inefficiency applied per step beyond efficient_step_threshold | |
| over_aggression choosing remove when ground truth is allow/flag | |
| """ | |
| from __future__ import annotations | |
| import os | |
| from dataclasses import dataclass, field | |
| from typing import Dict, Optional | |
| import yaml | |
| REWARD_CONFIG_PATH = os.getenv("REWARD_CONFIG", "reward_config.yaml") | |
| # --------------------------------------------------------------------------- | |
| # Default constants (used when YAML is absent) | |
| # --------------------------------------------------------------------------- | |
| _DEFAULTS: Dict[str, float | int] = { | |
| # Process | |
| "verify_claim_reward": 0.20, | |
| "check_source_reward": 0.20, | |
| "multi_source_bonus": 0.20, | |
| "multi_source_threshold": 3, # int: sources needed for bonus | |
| # Decision | |
| "decision_correct": 1.00, | |
| "decision_partial": 0.50, | |
| "decision_wrong": 0.00, | |
| # Penalties (negative values) | |
| "premature_decision": -0.50, | |
| "redundant_action": -0.10, | |
| "step_inefficiency": -0.05, | |
| "over_aggression": -0.30, | |
| # Clipping | |
| "min_total_reward": -2.00, | |
| "max_total_reward": 2.00, | |
| # Evidence gate | |
| "min_sources_for_decision": 2, # int | |
| "min_claims_verified_frac": 0.50, # fraction | |
| # Efficiency gate | |
| "efficient_step_threshold": 6, # int: steps before inefficiency kicks in | |
| } | |
| # Decision adjacency: maps (predicted, ground_truth) → partial credit | |
| _PARTIAL_CREDIT = { | |
| ("flag", "remove"): True, | |
| ("remove", "flag"): True, | |
| ("flag", "allow"): True, | |
| ("allow", "flag"): True, | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Config loader | |
| # --------------------------------------------------------------------------- | |
| class RewardConfig: | |
| def __init__(self, path: str = REWARD_CONFIG_PATH) -> None: | |
| self._cfg = dict(_DEFAULTS) | |
| if os.path.exists(path): | |
| with open(path, "r", encoding="utf-8") as fh: | |
| loaded = yaml.safe_load(fh) or {} | |
| self._cfg.update(loaded) | |
| def __getattr__(self, name: str): | |
| # Guard: _cfg may not exist yet during unpickling/__init__ | |
| if name == "_cfg": | |
| raise AttributeError(name) | |
| try: | |
| return self._cfg[name] | |
| except KeyError: | |
| raise AttributeError(f"RewardConfig has no attribute '{name}'") | |
| def as_dict(self) -> dict: | |
| return dict(self._cfg) | |
| # --------------------------------------------------------------------------- | |
| # Episode reward tracker | |
| # --------------------------------------------------------------------------- | |
| class EpisodeRewardState: | |
| """ | |
| Mutable accumulator attached to one environment episode. | |
| Tracks which claims/sources have already received process rewards. | |
| """ | |
| verified_claims: set[str] = field(default_factory=set) | |
| checked_sources: set[str] = field(default_factory=set) | |
| cumulative_reward: float = 0.0 | |
| step_count: int = 0 | |
| # Itemised totals | |
| verify_claim_total: float = 0.0 | |
| check_source_total: float = 0.0 | |
| multi_source_bonus_total: float = 0.0 | |
| decision_reward_total: float = 0.0 | |
| premature_penalty_total: float = 0.0 | |
| redundant_penalty_total: float = 0.0 | |
| inefficiency_penalty_total: float = 0.0 | |
| aggression_penalty_total: float = 0.0 | |
| def reset(self) -> None: | |
| self.verified_claims = set() | |
| self.checked_sources = set() | |
| self.cumulative_reward = 0.0 | |
| self.step_count = 0 | |
| self.verify_claim_total = 0.0 | |
| self.check_source_total = 0.0 | |
| self.multi_source_bonus_total = 0.0 | |
| self.decision_reward_total = 0.0 | |
| self.premature_penalty_total = 0.0 | |
| self.redundant_penalty_total = 0.0 | |
| self.inefficiency_penalty_total = 0.0 | |
| self.aggression_penalty_total = 0.0 | |
| # --------------------------------------------------------------------------- | |
| # Reward function | |
| # --------------------------------------------------------------------------- | |
| class RewardFunction: | |
| """ | |
| Stateless reward computer — receives episode state + action result, | |
| returns itemised reward dict and a human-readable message. | |
| """ | |
| def __init__(self, config: Optional[RewardConfig] = None) -> None: | |
| self._cfg = config or RewardConfig() | |
| # ------------------------------------------------------------------ | |
| # Public interface | |
| # ------------------------------------------------------------------ | |
| def compute_verify_claim( | |
| self, | |
| claim_id: str, | |
| state: EpisodeRewardState, | |
| ) -> tuple[float, str]: | |
| """Reward for verifying a claim.""" | |
| if claim_id in state.verified_claims: | |
| penalty = self._cfg.redundant_action | |
| state.redundant_penalty_total += penalty | |
| state.cumulative_reward += penalty | |
| return penalty, f"Redundant verification of {claim_id} ({penalty:+.2f})" | |
| reward = self._cfg.verify_claim_reward | |
| state.verified_claims.add(claim_id) | |
| state.verify_claim_total += reward | |
| state.cumulative_reward += reward | |
| # Step inefficiency | |
| ineff, msg = self._inefficiency(state) | |
| return reward + ineff, f"Verified claim {claim_id} ({reward:+.2f}){msg}" | |
| def compute_check_source( | |
| self, | |
| source_id: str, | |
| state: EpisodeRewardState, | |
| ) -> tuple[float, str]: | |
| """Reward for checking a source.""" | |
| if source_id in state.checked_sources: | |
| penalty = self._cfg.redundant_action | |
| state.redundant_penalty_total += penalty | |
| state.cumulative_reward += penalty | |
| return penalty, f"Redundant source check {source_id} ({penalty:+.2f})" | |
| reward = self._cfg.check_source_reward | |
| state.checked_sources.add(source_id) | |
| state.check_source_total += reward | |
| state.cumulative_reward += reward | |
| # Multi-source bonus | |
| bonus = 0.0 | |
| n_sources = len(state.checked_sources) | |
| threshold = int(self._cfg.multi_source_threshold) | |
| if n_sources == threshold: | |
| bonus = self._cfg.multi_source_bonus | |
| state.multi_source_bonus_total += bonus | |
| state.cumulative_reward += bonus | |
| # Step inefficiency | |
| ineff, msg = self._inefficiency(state) | |
| total = reward + bonus + ineff | |
| bonus_msg = f" +multi-source bonus ({bonus:+.2f})" if bonus else "" | |
| return total, f"Checked source {source_id} ({reward:+.2f}){bonus_msg}{msg}" | |
| def compute_decision( | |
| self, | |
| decision: str, | |
| ground_truth: str, | |
| state: EpisodeRewardState, | |
| num_claims: int, | |
| ) -> tuple[float, str]: | |
| """Reward/penalty for the final decision action.""" | |
| total = 0.0 | |
| messages = [] | |
| # ---- Premature decision gate ---- | |
| n_sources = len(state.checked_sources) | |
| n_verified = len(state.verified_claims) | |
| min_sources = int(self._cfg.min_sources_for_decision) | |
| min_verified = self._cfg.min_claims_verified_frac | |
| is_premature = ( | |
| n_sources < min_sources | |
| or (num_claims > 0 and n_verified / num_claims < min_verified) | |
| ) | |
| if is_premature: | |
| penalty = self._cfg.premature_decision | |
| state.premature_penalty_total += penalty | |
| state.cumulative_reward += penalty | |
| total += penalty | |
| messages.append(f"premature decision ({penalty:+.2f})") | |
| # ---- Correctness ---- | |
| if decision == ground_truth: | |
| reward = self._cfg.decision_correct | |
| messages.append(f"correct decision ({reward:+.2f})") | |
| elif _PARTIAL_CREDIT.get((decision, ground_truth), False): | |
| reward = self._cfg.decision_partial | |
| messages.append(f"partially correct decision ({reward:+.2f})") | |
| else: | |
| reward = self._cfg.decision_wrong | |
| messages.append("wrong decision (0.00)") | |
| state.decision_reward_total += reward | |
| state.cumulative_reward += reward | |
| total += reward | |
| # ---- Over-aggression penalty ---- | |
| if decision == "remove" and ground_truth in ("allow", "flag"): | |
| penalty = self._cfg.over_aggression | |
| state.aggression_penalty_total += penalty | |
| state.cumulative_reward += penalty | |
| total += penalty | |
| messages.append(f"over-aggressive moderation ({penalty:+.2f})") | |
| return total, "; ".join(messages) | |
| def final_score(self, state: EpisodeRewardState) -> float: | |
| """Return cumulative reward clipped to [min, max].""" | |
| return float(max( | |
| self._cfg.min_total_reward, | |
| min(self._cfg.max_total_reward, state.cumulative_reward), | |
| )) | |
| def breakdown_dict(self, state: EpisodeRewardState) -> dict: | |
| return { | |
| "verify_claim_reward": round(state.verify_claim_total, 4), | |
| "check_source_reward": round(state.check_source_total, 4), | |
| "multi_source_bonus": round(state.multi_source_bonus_total, 4), | |
| "decision_reward": round(state.decision_reward_total, 4), | |
| "premature_penalty": round(state.premature_penalty_total, 4), | |
| "redundant_penalty": round(state.redundant_penalty_total, 4), | |
| "inefficiency_penalty": round(state.inefficiency_penalty_total, 4), | |
| "aggression_penalty": round(state.aggression_penalty_total, 4), | |
| "cumulative": round(state.cumulative_reward, 4), | |
| "final_clipped": round(self.final_score(state), 4), | |
| } | |
| # ------------------------------------------------------------------ | |
| # Internal helpers | |
| # ------------------------------------------------------------------ | |
| def _inefficiency(self, state: EpisodeRewardState) -> tuple[float, str]: | |
| threshold = int(self._cfg.efficient_step_threshold) | |
| if state.step_count > threshold: | |
| penalty = self._cfg.step_inefficiency | |
| state.inefficiency_penalty_total += penalty | |
| state.cumulative_reward += penalty | |
| return penalty, f" inefficiency ({penalty:+.2f})" | |
| return 0.0, "" | |