""" Reward channels for E0 / E1 / E2, built for TRL's `normalize_then_sum`. AGGREGATION ----------- GRPOConfig(multi_objective_aggregation="normalize_then_sum") z-scores every reward function WITHIN its group before weighting and summing. That is exactly the decoupled GDPO-style normalization the brief asks for, so we get it without patching TRL. It also changes what the weights mean: alpha and gamma multiply STANDARDIZED channels, so alpha=0.5 reads as "half a standard deviation of diversity credit per standard deviation of quality". No manual rescaling of d_i into the 0-10 judge range is needed or wanted. GATING (why it is not simply "set it to 0") ------------------------------------------- The brief says a gate failure should zero the total reward. Under per-channel z-scoring, writing 0.0 into a channel does NOT mean "no credit" -- it means "whatever 0.0 ranks as inside this group". That is fine for quality (range [0,10], floor 0) and for deviation (range [0,2], floor 0), but it is actively WRONG for the marginal contribution: m_i = logdet(L) - logdet(L_-i) <= log(1+eps) ~ 0 m_i is always <= 0, so 0.0 is its CEILING. Gating a broken story to 0.0 in the marginal channel would hand it the single highest diversity credit in the group -- a reward-hacking channel we would have built ourselves. So the uniform rule, applied to every diversity channel regardless of sign convention: an ineligible sample is assigned the MINIMUM value among eligible samples in its group. It can never out-rank a sample that earned its credit. Combined with quality -> 0.0 (a hard floor in that channel), a gate-failed story lands at the bottom of every channel it participates in. CONSTANT-REWARD SAFETY ---------------------- If every sample in a group is gated, a channel goes constant; TRL's (x - mean)/(std + 1e-4) then yields ~0 for all of them. That is the correct outcome: a group with no valid samples carries no signal. It is not a crash and not a NaN, but it IS worth logging, so `frac_groups_degenerate` is tracked. """ from __future__ import annotations import statistics from dataclasses import dataclass, field import numpy as np import gates as G from diversity import l2_normalize, marginal_contributions, pairwise_deviation, zscore # --------------------------------------------------------------- embeddings _ENCODER = None _EMB_MODEL = "BAAI/bge-base-en-v1.5" def get_encoder(model_name: str = _EMB_MODEL, device: str | None = None): """Module-level singleton; ~110M params (~0.22GB), negligible next to the policy. Device is overridable via EMB_DEVICE so tests (and any process running alongside a training job that already owns the GPU) can force CPU. """ global _ENCODER if _ENCODER is None: import os from sentence_transformers import SentenceTransformer dev = device or os.environ.get("EMB_DEVICE", "cuda") _ENCODER = SentenceTransformer(model_name, device=dev) return _ENCODER def embed(texts: list[str]) -> np.ndarray: if not texts: return np.zeros((0, 768)) E = get_encoder().encode( texts, normalize_embeddings=True, batch_size=32, show_progress_bar=False, convert_to_numpy=True, ) return l2_normalize(np.asarray(E, dtype=np.float64)) # ------------------------------------------------------------------ config @dataclass class RewardConfig: arm: str = "E0" # E0 | E1 | E2 alpha: float = 0.5 # weight on deviation channel gamma: float = 0.5 # weight on marginal channel tau: float = 5.0 # quality gate for diversity credit min_words: int = G.MIN_WORDS max_words: int = G.MAX_WORDS def channels(self) -> list[str]: if self.arm == "E0": return ["quality"] if self.arm == "E1": return ["quality", "deviation"] if self.arm == "E2": return ["quality", "deviation", "marginal"] raise ValueError(self.arm) def weights(self) -> list[float]: return {"E0": [1.0], "E1": [1.0, self.alpha], "E2": [1.0, self.alpha, self.gamma]}[self.arm] @dataclass class StepStats: n: int = 0 gate_pass: float = 0.0 ends_cleanly: float = 0.0 mean_quality: float = 0.0 mean_quality_passing: float = 0.0 mean_novelty: float = 0.0 frac_above_tau: float = 0.0 mean_deviation: float = 0.0 mean_logdet: float = 0.0 mean_marginal: float = 0.0 mean_words: float = 0.0 frac_groups_degenerate: float = 0.0 reasons: dict = field(default_factory=dict) def _gate_floor(values: np.ndarray, eligible: np.ndarray) -> np.ndarray: """Ineligible samples take the minimum value among eligible ones. Sign-convention agnostic: works for deviation (>=0) and for marginal (<=0) alike. If nothing is eligible, the channel is flat -> z-scores to 0 in TRL. """ out = values.astype(np.float64).copy() if not eligible.any(): return np.zeros_like(out) out[~eligible] = values[eligible].min() return out class RewardEngine: """Scores one GRPO batch: gates -> judge -> embeddings -> per-channel values. TRL calls each reward function separately, but we want ONE judge call set and ONE embedding pass per batch. So the engine computes everything once and memoizes on the batch signature; the per-channel closures just read it. """ def __init__(self, cfg: RewardConfig, judge, wandb_run=None, log_prefix="train"): self.cfg = cfg self.judge = judge self.wandb_run = wandb_run self.log_prefix = log_prefix self._sig = None self._cache: dict[str, np.ndarray] = {} self.last_stats: StepStats | None = None self.history: list[StepStats] = [] # ---- core ---------------------------------------------------------- def compute(self, prompts: list[str], texts: list[str], finish_reasons: list[str] | None = None) -> dict[str, np.ndarray]: sig = hash((tuple(prompts), tuple(texts))) if sig == self._sig: return self._cache n = len(texts) finish_reasons = finish_reasons or [None] * n # 1. programmatic gates (free, run first) gres = [G.check(t, finish_reason=fr, min_words=self.cfg.min_words, max_words=self.cfg.max_words) for t, fr in zip(texts, finish_reasons)] passed = np.array([r.passed for r in gres], dtype=bool) # 2. judge only the stories that survived the gates -- never pay to # score text we have already decided to zero out. quality = np.zeros(n); novelty = np.zeros(n) idx = [i for i in range(n) if passed[i]] if idx: scores = self.judge.score_many_sync([(prompts[i], texts[i]) for i in idx]) for i, s in zip(idx, scores): quality[i] = s.quality novelty[i] = s.novelty # 3. embeddings + per-group diversity E = embed(texts) groups: dict[str, list[int]] = {} for i, p in enumerate(prompts): groups.setdefault(p, []).append(i) dev = np.zeros(n); marg = np.zeros(n) logdets, degenerate = [], 0 for _, ids in groups.items(): sub = E[ids] d = pairwise_deviation(sub) m = marginal_contributions(sub) # z-score m within group: raw m has a long negative tail (a duplicate # pair reaches log(eps) ~ -6.9) that would otherwise dominate. mz = zscore(m) for k, i in enumerate(ids): dev[i] = d[k]; marg[i] = mz[k] from diversity import logdet_volume logdets.append(logdet_volume(sub)) if not passed[ids].any(): degenerate += 1 # 4. eligibility for diversity credit: gates AND quality >= tau. # Conditioning is what stops "incoherent but different" from paying. eligible = passed & (quality >= self.cfg.tau) dev_c = np.zeros(n); marg_c = np.zeros(n) for _, ids in groups.items(): ids_a = np.array(ids) dev_c[ids_a] = _gate_floor(dev[ids_a], eligible[ids_a]) marg_c[ids_a] = _gate_floor(marg[ids_a], eligible[ids_a]) quality_c = np.where(passed, quality, 0.0) out = {"quality": quality_c, "deviation": dev_c, "marginal": marg_c} self._sig, self._cache = sig, out # ---- stats ---- from collections import Counter cnt = Counter(r for x in gres for r in x.reasons) st = StepStats( n=n, gate_pass=float(passed.mean()), ends_cleanly=float(np.mean([r.completeness for r in gres])), mean_quality=float(quality.mean()), mean_quality_passing=float(quality[passed].mean()) if passed.any() else 0.0, mean_novelty=float(novelty[passed].mean()) if passed.any() else 0.0, frac_above_tau=float(eligible.mean()), mean_deviation=float(dev.mean()), mean_logdet=float(np.mean(logdets)) if logdets else 0.0, mean_marginal=float(marg.mean()), mean_words=float(np.mean([r.n_words for r in gres])), frac_groups_degenerate=degenerate / max(1, len(groups)), reasons=dict(cnt), ) self.last_stats = st self.history.append(st) self._log(st) return out def _log(self, st: StepStats) -> None: print(f" [rw] gate={st.gate_pass:.2f} end={st.ends_cleanly:.2f} " f"q={st.mean_quality_passing:.2f} >tau={st.frac_above_tau:.2f} " f"dev={st.mean_deviation:.3f} logdet={st.mean_logdet:.2f} " f"w={st.mean_words:.0f} {st.reasons if st.reasons else ''}", flush=True) if self.wandb_run is not None: p = self.log_prefix self.wandb_run.log({ f"{p}/gate_pass": st.gate_pass, f"{p}/ends_cleanly": st.ends_cleanly, f"{p}/quality_passing": st.mean_quality_passing, f"{p}/quality_all": st.mean_quality, f"{p}/novelty": st.mean_novelty, f"{p}/frac_above_tau": st.frac_above_tau, f"{p}/deviation": st.mean_deviation, f"{p}/logdet": st.mean_logdet, f"{p}/marginal_z": st.mean_marginal, f"{p}/words": st.mean_words, f"{p}/groups_degenerate": st.frac_groups_degenerate, }) # ---- TRL adapters --------------------------------------------------- def make_reward_funcs(self): """Return TRL-compatible reward callables, one per active channel.""" funcs = [] for ch in self.cfg.channels(): funcs.append(self._make(ch)) return funcs def _make(self, channel: str): engine = self def f(completions, prompts=None, **kwargs): texts = [_text(c) for c in completions] ps = [_ptext(p) for p in (prompts or [""] * len(texts))] fr = kwargs.get("finish_reasons") vals = engine.compute(ps, texts, fr) return [float(x) for x in vals[channel]] f.__name__ = f"{channel}_reward" return f def _text(c) -> str: if isinstance(c, list): return c[-1].get("content", "") if c else "" if isinstance(c, dict): return c.get("content", "") return str(c) def _ptext(p) -> str: if isinstance(p, list): # chat format: the user turn carries the writing prompt for m in reversed(p): if isinstance(m, dict) and m.get("role") == "user": return m.get("content", "") return p[-1].get("content", "") if p else "" return str(p)