""" TICKETMELT — Reward functions. Four independent reward components are computed at the END of each episode (not per-step). This is cleaner for GRPO because GRPO ranks whole completions. The continuous components are logged for plots and analysis. The binary wrapper is what gets returned to TRL for training. Design principle: each component targets a different behavior. An agent can't game one without leaving easy wins on the others. """ from __future__ import annotations from dataclasses import dataclass from .models import State, RoundRecord # --------------------------------------------------------------------------- # Reward components (continuous, for logging and analysis) # --------------------------------------------------------------------------- def reward_r1_service_restored(state: State) -> float: """ R1: Did MY service (the trained engineer's service) recover before deadline? 1.0 = recovered before deadline 0.3 = recovered but after deadline (partial credit) 0.0 = didn't recover """ my_service = state.services[state.trained_engineer] if not my_service.completed: return 0.0 if my_service.completed_on_time: return 1.0 return 0.3 def reward_r2_site_uptime(state: State) -> float: """ R2: What fraction of all services recovered on time? Rewards cooperative play — I can't win alone by ignoring peers. """ total = len(state.services) on_time = sum(1 for s in state.services.values() if s.completed_on_time) return on_time / total if total > 0 else 0.0 def reward_r3_clean_deploys(state: State) -> float: """ R3: How often did I avoid collisions when I deployed? Only counts rounds where I DID deploy (not MONITOR rounds). A model that always monitors gets 0.5 (neutral) rather than 1.0, so it can't win R3 by doing nothing. """ me = state.trained_engineer my_deploy_rounds = [] my_collision_rounds = [] for record in state.history: commitment = record.commitments.get(me) if commitment in ("DEPLOY_PROD_A", "DEPLOY_PROD_B"): my_deploy_rounds.append(record) server_key = f"PROD_{commitment[-1]}" if server_key in record.collisions: my_collision_rounds.append(record) if not my_deploy_rounds: return 0.0 # never deployed = no credit collision_rate = len(my_collision_rounds) / len(my_deploy_rounds) return 1.0 - collision_rate def reward_r4_yield_to_critical(state: State) -> float: """ R4: Did I yield (MONITOR) when a peer signaled urgency AND had more work remaining? This is the "adaptation to heterogeneous urgency" signal. The urgency_flag is private to each engineer, but the scripted peers never actually signal urgency explicitly — we detect urgency structurally by checking who had a tighter (deadline - fix_rounds_remaining) margin. Score = fraction of "yield-worthy" rounds where I correctly yielded. """ me = state.trained_engineer my_service = state.services[me] yield_opportunities = 0 correct_yields = 0 for i, record in enumerate(state.history): my_commitment = record.commitments.get(me) # A "yield-worthy" round is one where at least one peer had a tighter # margin than me AND was about to contend for a server I wanted. # Simplification: a peer is "critical" if their slack <= 1 # (slack = deadline - current_round - fix_rounds_remaining) # And I'm "less critical" if my slack > their slack. my_slack = my_service.deadline_round - (i + 1) - _rounds_remaining_at(state, me, i) for peer_name, peer_service in state.services.items(): if peer_name == me: continue peer_slack = peer_service.deadline_round - (i + 1) - _rounds_remaining_at(state, peer_name, i) peer_wanted_server = record.commitments.get(peer_name) in ("DEPLOY_PROD_A", "DEPLOY_PROD_B") if peer_slack <= 1 and my_slack > peer_slack and peer_wanted_server: yield_opportunities += 1 if my_commitment == "MONITOR": correct_yields += 1 break # one opportunity per round max if yield_opportunities == 0: return 0.5 # no opportunities arose, neutral score return correct_yields / yield_opportunities def _rounds_remaining_at(state: State, engineer: str, round_idx: int) -> int: """Helper: how many fix-rounds did `engineer` have remaining AT round_idx?""" # Reconstruct from history. fix_rounds_remaining decrements only on successful solo deploy. service = state.services[engineer] total = service.fix_rounds_total successful_deploys_before = sum( 1 for r in state.history[:round_idx + 1] if r.successful_deploys.get(engineer) is not None ) return max(0, total - successful_deploys_before) # --------------------------------------------------------------------------- # Composite scoring # --------------------------------------------------------------------------- @dataclass class RewardBreakdown: r1: float r2: float r3: float r4: float weighted_sum: float binary_grpo_reward: float DEFAULT_WEIGHTS = {"r1": 0.5, "r2": 0.3, "r3": 0.1, "r4": 0.1} DEFAULT_GOODNESS_THRESHOLD = 0.45 # raised from 0.55 after observing inflated R3 def compute_rewards( state: State, weights: dict[str, float] = None, goodness_threshold: float = DEFAULT_GOODNESS_THRESHOLD, ) -> RewardBreakdown: """ Compute all four reward components plus the composite used for training. The `binary_grpo_reward` is what TRL sees. The continuous fields are for logging and plots. """ if weights is None: weights = DEFAULT_WEIGHTS r1 = reward_r1_service_restored(state) r2 = reward_r2_site_uptime(state) r3 = reward_r3_clean_deploys(state) r4 = reward_r4_yield_to_critical(state) weighted = weights["r1"] * r1 + weights["r2"] * r2 + weights["r3"] * r3 + weights["r4"] * r4 binary = 1.0 if weighted > goodness_threshold else 0.0 return RewardBreakdown( r1=r1, r2=r2, r3=r3, r4=r4, weighted_sum=weighted, binary_grpo_reward=binary, )