training/rewards.py

"""Runbook-compliance reward for PM-Ops GRPO training.

Design philosophy (v4)
----------------------
Previous reward formulas (json_ratio, ok_ratio, diversity_bonus) were all constant
once the model learned to output valid JSON — giving zero GRPO advantage.

Root problem: env_score was always 0 because:
  1. The model used hardcoded label='sec-issue' / priority='P1', which fail
     TicketingApp validation for most org configs → no ticket created → grader
     returns 0.0 immediately.
  2. Dataset seeds produced triage briefs but the env ran a DIFFERENT task type
     (release_notes, dep_update) for the same seed → guaranteed mismatch.

Fix: reward is computed from the agent's ACTIONS compared to RUNBOOK DATA.
The org_config (returned by meta.read_runbook) varies by seed — different orgs
have different valid labels, priorities, teams, and oncall channels. The model's
fixed template (sec-issue, P1, infra) scores well for some orgs and badly for
others, creating the reward VARIANCE that GRPO needs.

Components (sum = 1.0 when all correct):
  read_runbook   0.10  — process: did agent read runbook first?
  valid_label    0.20  — used a label from label_taxonomy? (+0.20 / -0.10)
  valid_priority 0.15  — used a priority from priority_levels? (+0.15 / -0.10)
  valid_team     0.20  — assigned to a team from team_map? (+0.20 / -0.10)
  right_channel  0.25  — posted to an oncall channel? (+0.25 / -0.10 per wrong)
  env_bonus      0.10  — env grader bonus when everything lines up correctly

This is computed IN the rollout (not by TRL reward_funcs) because it needs
access to the per-episode runbook data.
"""


def compute_rollout_reward(
    *,
    read_runbook_done: bool,
    valid_labels: set,        # from org_config.label_taxonomy.values()
    valid_priorities: set,    # from org_config.priority_levels
    valid_teams: set,         # from org_config.team_map.values()
    oncall_channels: set,     # from org_config.oncall_channels.values()
    ticket_label: str | None,       # label used in create_ticket (None if no ticket)
    ticket_priority: str | None,    # priority used in create_ticket
    assigned_team: str | None,      # team from assign_ticket (None if not called)
    posted_channels: list[str],     # all channels from chat.post_message
    env_score: float,               # final env grader score (0–1)
    valid_json_count: int,          # steps with parseable JSON output
) -> float:
    """Compute a single reward scalar for one episode.

    Returns a value in [-1.0, 1.0].
    """
    if valid_json_count == 0:
        return -1.0

    reward = 0.0

    # 1. Runbook read (+0.10 process bonus)
    if read_runbook_done:
        reward += 0.10

    # 2. Ticket label valid (only scored if a ticket was created)
    if ticket_label is not None:
        if valid_labels:
            if ticket_label in valid_labels:
                reward += 0.20
            else:
                reward -= 0.10  # wrong label — validation would have rejected it

    # 3. Ticket priority valid
    if ticket_priority is not None:
        if valid_priorities:
            if ticket_priority in valid_priorities:
                reward += 0.15
            else:
                reward -= 0.10

    # 4. Ticket assigned to a valid team
    if assigned_team is not None:
        if valid_teams:
            if assigned_team in valid_teams:
                reward += 0.20
            else:
                reward -= 0.10

    # 5. Posted to the right oncall channel
    if posted_channels:
        if oncall_channels:
            correct_posts = [ch for ch in posted_channels if ch in oncall_channels]
            wrong_posts   = [ch for ch in posted_channels if ch not in oncall_channels]
            if correct_posts:
                reward += 0.25
            reward -= 0.10 * len(wrong_posts)  # -0.10 per channel-spray post
        else:
            # Posted without reading runbook — can't verify, mild penalty
            reward -= 0.05 * len(posted_channels)

    # 6. Env grader bonus — full env score adds on top when everything is correct
    reward += env_score * 0.10

    return max(-1.0, min(1.0, reward))