""" 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 # --------------------------------------------------------------------------- @dataclass 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, ""