training/rewards.py

"""
permanence.training.rewards — composable reward functions for GRPO.

The reward stack has two complete components, separated by the source of
information they can see:

    1. Text-only reward (``reward_format``) — inspects the completion
       string only. Handles tag compliance and brevity together; both are
       functions of the text alone.

    2. Environmental reward — steps the env with each completion and
       returns the env's scalar reward. Wired in
       ``training/stages/stage_3_grpo.py`` because it needs a live env
       handle; wrapped here by ``weighted_environmental_reward`` to add
       dynamic scheduling.

Dynamic scheduling phases the format reward out as the environmental
reward takes over: the model starts with strong pressure to produce
well-formed tags, then pressure shifts to predicting correctly.

Two optional shaping knobs are exposed but disabled by default
(``BETA_RANK = 0.0``). They are retained in code only because small
values can still be useful for ablation; see the block comment below.
"""
from __future__ import annotations

import re
from collections import deque
from dataclasses import dataclass, field
from typing import Callable, Deque, Dict, List, Optional


ACTION_TAG_RE = re.compile(r"<action\s+id=[\"'][^\"']+[\"']", re.IGNORECASE)
REVERSIBILITY_TAG_RE = re.compile(r"<reversibility\s+level=[\"'][Rr][1-5][\"']", re.IGNORECASE)
LEVEL_RE = re.compile(r"level=[\"']([Rr])([1-5])[\"']", re.IGNORECASE)
CONFIDENCE_RE = re.compile(r"confidence=[\"']([0-9.]+)[\"']", re.IGNORECASE)
THINKING_RE = re.compile(r"<thinking>.*?</thinking>", re.IGNORECASE | re.DOTALL)


# ─────────────────────────────────────────────────────────────────────────────
# Shaping knobs
# ─────────────────────────────────────────────────────────────────────────────
#
# Rank-based "unlikeliness" shaping (He et al., arXiv:2506.02355) was
# designed for *binary-verifier* RL tasks (a proof either checks or it
# doesn't). Classification-style RLVR with continuous partial-credit
# rewards — the setting here — does not benefit from rank-based shaping:
# when the correct prediction also earns the highest raw reward, a
# penalty on the top-ranked sample inverts the gradient signal and the
# policy drifts toward the *wrong* answer. We ship with
# ``BETA_RANK = 0.0`` and keep the plumbing only so the effect can be
# re-measured via an explicit ablation.
#
# The R-level balance bonus is a small additive bonus (+0.1 × (R − 3))
# applied when a correct prediction lands on the rarer high-R classes
# (R4 and R5). Its purpose is to counteract the base-rate imbalance in
# the training distribution; the bonus is conservative enough that it
# cannot, on its own, flip the gradient direction.

BETA_RANK = 0.0                  # disabled by default; see block comment above
R_LEVEL_BONUS_PER_LEVEL = 0.1    # additive bonus per R-level on correct R4/R5
#
#
# Research basis: The unlikeliness-reward technique in He et al. was designed
# for FORMAL THEOREM PROVING with BINARY rewards (proof works / doesn't).
# Our task is a classification-style RLVR with CONTINUOUS partial-credit
# rewards (level_accuracy × calibration in [0, 1]). Applying unlikeliness
# to our continuous-reward setting has the opposite of the intended effect:
# it penalizes correct, confident predictions (high reward) relative to
# wrong-but-close predictions (lower but still positive reward).
#
# Empirical evidence from a pilot run: the wrong prediction (R1 on an
# actual R2 action) collected a higher mean reward than the correct one
# because the rank-based penalty on the top-ranked sample bit into the
# correct prediction more than the wrong one.
#   The unlikeliness shaping inverted the gradient: R1 paid MORE than R2.
#   GRPO learned to predict R1 and eval accuracy dropped to 46%.
#
# Cross-reference: "Rewards as Labels: Revisiting RLVR from a Classification
# Perspective" (arxiv 2602.05630) identifies GRPO's "Gradient Misassignment
# in Positives" for classification tasks. Unlikeliness shaping amplifies
# this pathology rather than fixing it in our setting.
#
# Setting BETA_RANK=0.0 disables the shaping entirely. The forced
# variants + R-level balance bonus prevent the degenerate "always-safe"
# policy without needing rank-based shaping. Our
# classification-style reward already has a clear gradient signal without
# needing unlikeliness to surface rare samples.
BETA_RANK = 0.0  # disabled — see note below
R_LEVEL_BONUS_PER_LEVEL = 0.1  # Additive bonus per R-level of correct rare prediction


# ─────────────────────────────────────────────────────────────────────────────
# Reward 1 — pure-text format + brevity
# ─────────────────────────────────────────────────────────────────────────────


def reward_format(completions: List[str], **_: object) -> List[float]:
    """Range: -0.1 (rambling, no tags) to +1.0 (perfect, concise).

    Partial-credit grid chosen so EVERY rollout earns a different value
    unless the group is literally identical — keeps ``reward_std > 0`` so
    GRPO has a gradient.

        +0.20  ``<action id="…"``
        +0.15  action tag is closed
        +0.15  ``<reversibility level="Rn"``
        +0.10  reversibility tag is closed
        +0.10  ``<thinking>…</thinking>`` block present
        +0.10  confidence attribute parses as a float in [0, 1]
        +0.20  length ≤ 400 chars (strongest brevity tier)
        +0.10  length 400–600 chars
         0.00  length 600–900 chars
        -0.10  length > 1100 chars (rambling — drift signal)

    Notice brevity is folded in directly: an earlier iteration had a separate
    ``reward_brevity`` callable returning a constant 0.15, which added
    noise to logging without real signal. Length belongs with
    format because both are string-only properties.
    """
    scores: List[float] = []
    for text in completions:
        s = 0.0
        if ACTION_TAG_RE.search(text):
            s += 0.20
            if "/>" in text:
                s += 0.15
        rev_match = REVERSIBILITY_TAG_RE.search(text)
        if rev_match:
            s += 0.15
            tail = text[rev_match.start():]
            if "/>" in tail or "</reversibility>" in tail:
                s += 0.10
        if THINKING_RE.search(text):
            s += 0.10
        conf_match = CONFIDENCE_RE.search(text)
        if conf_match:
            try:
                c = float(conf_match.group(1))
                if 0.0 <= c <= 1.0:
                    s += 0.10
            except (TypeError, ValueError):
                pass
        n = len(text)
        if n <= 400:
            s += 0.20
        elif n <= 600:
            s += 0.10
        elif n > 1100:
            s -= 0.10
        scores.append(max(-0.10, min(1.0, s)))
    return scores


# ─────────────────────────────────────────────────────────────────────────────
# Dynamic reward weighting + length monitoring
# ─────────────────────────────────────────────────────────────────────────────


@dataclass
class RewardSchedule:
    """Piecewise-linear weight schedule across training.

    Format dominates in the first 60 steps (training wheels), phases out by
    step 150 so the environmental reward (which carries the actual task
    signal) takes over.
    """

    total_episodes: int = 300

    def weight_format(self, episode: int) -> float:
        if episode < 60:
            return 1.0 - 0.8 * (episode / 60)
        if episode < 150:
            return 0.2 * (1.0 - (episode - 60) / 90)
        return 0.0

    def weight_environmental(self, episode: int) -> float:
        """Env reward is the workhorse. Starts at 0.5 (while format trains
        the model to produce valid output) and ramps to 1.5 by step 150."""
        if episode < 60:
            return 0.5 + 0.5 * (episode / 60)
        if episode < 150:
            return 1.0 + 0.5 * ((episode - 60) / 90)
        return 1.5

    def weights_at(self, episode: int) -> List[float]:
        return [self.weight_format(episode), self.weight_environmental(episode)]


@dataclass
class LengthMonitor:
    """Rolling-average length tracker with an abort flag.

    When the mean of the last ``window`` completion lengths exceeds
    ``threshold_chars`` for ``trigger_windows`` consecutive windows, sets
    ``abort_flag=True``. Stage 3 checks this before each GRPO step and
    raises a clean abort error.
    """

    window: int = 20
    threshold_chars: int = 1000
    trigger_windows: int = 3
    recent_lengths: Deque[int] = field(default_factory=lambda: deque(maxlen=20))
    consecutive_over: int = 0
    abort_flag: bool = False

    def observe(self, completion: str) -> None:
        self.recent_lengths.append(len(completion))
        if len(self.recent_lengths) < self.window:
            return
        avg = sum(self.recent_lengths) / len(self.recent_lengths)
        if avg > self.threshold_chars:
            self.consecutive_over += 1
        else:
            self.consecutive_over = 0
        if self.consecutive_over >= self.trigger_windows:
            self.abort_flag = True


# ─────────────────────────────────────────────────────────────────────────────
# Reward-pack builder
# ─────────────────────────────────────────────────────────────────────────────


@dataclass
class RewardPack:
    """Container for the two weighted reward callables plus the shared
    episode counter and length monitor.

    The environmental reward is NOT in ``funcs`` because it needs access
    to the training log (side effect). Stage 3 constructs it separately
    and appends it to the list before giving it to the GRPO trainer.
    """

    funcs: List[Callable[..., List[float]]]
    schedule: RewardSchedule
    length_monitor: LengthMonitor
    episode_counter: List[int] = field(default_factory=lambda: [0])


def build_reward_pack(total_episodes: int = 300) -> RewardPack:
    """Assemble the text-only reward pack.

    Stage 3 pairs this with a separately-constructed environmental reward
    function that runs env.step internally. The two rewards together form
    the complete signal.
    """
    schedule = RewardSchedule(total_episodes=total_episodes)
    monitor = LengthMonitor()
    ep_counter = [0]

    def make_weighted(fn: Callable[..., List[float]], weight_fn: Callable[[int], float]) -> Callable[..., List[float]]:
        def wrapped(completions: List[str] | None = None, **kwargs) -> List[float]:
            # Handle completions-as-positional-or-kwarg so TRL's
            # ``prompts=..., completions=...`` calling convention doesn't
            # cause an arg-conflict when forwarding to inner functions.
            if completions is None:
                completions = kwargs.pop("completions", [])
            for c in completions:
                monitor.observe(c)
            w = weight_fn(ep_counter[0])
            if w == 0.0:
                return [0.0] * len(completions)
            # ``reward_format`` accepts ``**_`` so it absorbs everything —
            # passing completions as a kwarg is safe and collision-free.
            raw = fn(completions=completions, **kwargs)
            return [w * r for r in raw]

        wrapped.__name__ = fn.__name__
        return wrapped

    funcs = [
        make_weighted(reward_format, schedule.weight_format),
    ]
    return RewardPack(funcs=funcs, schedule=schedule, length_monitor=monitor, episode_counter=ep_counter)


def weighted_environmental_reward(
    raw_fn: Callable[..., List[float]],
    pack: RewardPack,
) -> Callable[..., List[float]]:
    """Wrap an environmental reward fn with three shaping steps:

    1. **Schedule weighting** — multiply by the current env weight from
       the pack's schedule (grows from 0.5 → 1.5 over 150 steps).

    2. **Unlikeliness reward** (He et al. 2506.02355) — within each group
       of rollouts, rank samples by raw reward. Apply a multiplicative
       penalty (1 - β_rank × rank_norm) to high-reward samples so rare
       low-reward-but-still-positive samples get stronger relative
       advantages. This breaks the "always pick the safe action" local
       optimum that a naive "prefer safe action" policy would find.

    3. **R-level balance bonus** — read the last training-log entry's
       (predicted_r_level, actual_r_level) pair; if the agent correctly
       predicted a rare high-R action (R4 or R5), add a small bonus.
       This directly incentivizes developing the R4/R5 prediction
       capability that the policy would otherwise underweight on base-rate grounds.

    The wrapped function forwards ALL kwargs straight through (without
    making completions a positional arg) so TRL's usual ``prompts=...``
    keyword does not collide with the wrapped function's positional
    ``prompts`` parameter. The pipeline previously crashed on exactly this
    bug — the fix is to forward every arg by keyword only.
    """

    def wrapped(completions: List[str] | None = None, **kwargs) -> List[float]:
        if completions is None:
            completions = kwargs.pop("completions", [])
        for c in completions:
            pack.length_monitor.observe(c)

        w = pack.schedule.weight_environmental(pack.episode_counter[0])
        if w == 0.0:
            return [0.0] * len(completions)

        # Step 1: raw env reward
        raw = raw_fn(completions=completions, **kwargs)

        # Step 2: unlikeliness reward shaping (He et al. 2025).
        # Rank samples in descending reward order; apply multiplicative
        # penalty (1 - β_rank × rank_norm) to high-reward samples so rare
        # low-reward successful samples get stronger relative advantages.
        #
        # Only apply to positive rewards — we never up-weight losses.
        G = len(raw)
        if G >= 2:
            sorted_indices = sorted(range(G), key=lambda i: -raw[i])
            rank_of = {idx: r for r, idx in enumerate(sorted_indices)}
            shaped = []
            for i in range(G):
                rank_norm = (G - 1 - rank_of[i]) / max(G, 1)
                if raw[i] > 0:
                    mult = 1.0 - BETA_RANK * rank_norm
                else:
                    mult = 1.0
                shaped.append(raw[i] * mult)
        else:
            shaped = list(raw)

        # Step 3: R-level balance bonus from the training log.
        # ``_make_task_reward`` exposes ``training_log`` on the returned
        # callable (see stage_3_grpo). The last G entries correspond to
        # the current batch of completions. Bonus for correctly predicting
        # R4 or R5 (the rare classes the policy avoids).
        training_log = getattr(raw_fn, "training_log", None)
        if training_log is not None and len(training_log) >= G:
            recent = training_log[-G:]
            for i, entry in enumerate(recent):
                pred = entry.get("predicted_r_level")
                actual = entry.get("action_r_level") or entry.get("actual_r_level")
                if pred is None or actual is None:
                    continue
                if pred == actual and actual >= 4:
                    shaped[i] += R_LEVEL_BONUS_PER_LEVEL * (actual - 3)

        return [w * r for r in shaped]

    wrapped.__name__ = raw_fn.__name__
    return wrapped