Blame for reward.py

"""
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, ""