""" Tests for the reward-channel construction in rewards.py, with a fake judge so nothing hits the network. The load-bearing test is test_gated_sample_never_outranks_in_marginal_channel: m_i <= log(1+eps) ~ 0 is ALWAYS negative, so gating a failed story to 0.0 would hand it the highest diversity credit in the group. That would be a reward-hacking channel we built ourselves, and it is exactly the kind of sign error that is invisible in aggregate training curves. """ import numpy as np import gates as G import rewards as R from judge import JudgeScore class FakeJudge: """Scores by a marker embedded in the text; no network.""" def __init__(self, mapping=None): self.mapping = mapping or {} self.calls = 0 def score_many_sync(self, pairs): self.calls += len(pairs) out = [] for _, story in pairs: q = 8.0 for k, v in self.mapping.items(): if k in story: q = v break out.append(JudgeScore(quality=q, novelty=5.0)) return out def _story(marker: str, n_words: int = 250, seed: int = 0) -> str: """Gate-passing filler with a marker and enough lexical variety to clear the 4-gram-loop and entropy gates.""" rng = np.random.default_rng(seed) vocab = ["harbor", "clock", "ember", "listen", "gravel", "orchard", "signal", "letter", "winter", "throat", "marble", "engine", "sister", "quiet", "amber", "hollow", "ribbon", "tunnel", "pepper", "静"][:19] words = [vocab[i] for i in rng.integers(0, len(vocab), n_words)] return f"{marker} " + " ".join(words) + "." def _engine(arm="E2", tau=5.0, judge=None): cfg = R.RewardConfig(arm=arm, tau=tau) return R.RewardEngine(cfg, judge or FakeJudge()) def _fake_embed(monkey_vals): """Patch rewards.embed to return a fixed matrix.""" R.embed = lambda texts: monkey_vals # ------------------------------------------------------------------ configs def test_channels_and_weights_per_arm(): assert R.RewardConfig(arm="E0").channels() == ["quality"] assert R.RewardConfig(arm="E1").channels() == ["quality", "deviation"] assert R.RewardConfig(arm="E2").channels() == ["quality", "deviation", "marginal"] assert R.RewardConfig(arm="E2", alpha=0.3, gamma=0.7).weights() == [1.0, 0.3, 0.7] # ------------------------------------------------------------------- gating def test_gate_failure_zeroes_quality_channel(): eng = _engine("E1") texts = [_story("GOOD", 250, i) for i in range(3)] + ["too short."] prompts = ["p"] * 4 out = eng.compute(prompts, texts) assert out["quality"][3] == 0.0, "gate-failed story must floor the quality channel" assert (out["quality"][:3] > 0).all() def test_judge_is_not_called_for_gate_failures(): """We must never pay to score text we have already decided to zero.""" j = FakeJudge() eng = _engine("E1", judge=j) texts = [_story("A", 250, 1), "nope.", "also short."] eng.compute(["p"] * 3, texts) assert j.calls == 1, f"judge called {j.calls} times, expected 1" def test_low_quality_forfeits_diversity_credit(): """tau conditioning: a coherent but low-quality story earns no diversity.""" j = FakeJudge({"BAD": 2.0, "GOOD": 8.0}) eng = _engine("E1", tau=5.0, judge=j) texts = [_story("GOOD", 250, i) for i in range(3)] + [_story("BAD", 250, 9)] out = eng.compute(["p"] * 4, texts) dev = out["deviation"] assert dev[3] <= dev[:3].min() + 1e-12, \ f"sub-tau story got deviation credit {dev[3]} vs eligible min {dev[:3].min()}" def test_gated_sample_never_outranks_in_marginal_channel(): """THE sign trap. m_i is always <= 0, so gating to 0.0 would make failure the single best value in the channel.""" j = FakeJudge({"BAD": 1.0, "GOOD": 8.0}) eng = _engine("E2", tau=5.0, judge=j) texts = [_story("GOOD", 250, i) for i in range(5)] + [_story("BAD", 250, 42)] out = eng.compute(["p"] * 6, texts) m = out["marginal"] assert m[5] <= m[:5].min() + 1e-12, \ f"ineligible story ranked ABOVE eligible ones in marginal channel: {m}" assert m[5] != 0.0 or np.allclose(m, 0.0), "suspicious exact-zero gate value" def test_all_gated_group_is_constant_not_nan(): """Every sample failing => channels go constant. TRL's (x-mean)/(std+1e-4) then yields ~0 for all, which is correct (no signal), and must not be NaN.""" eng = _engine("E2") texts = ["short."] * 4 out = eng.compute(["p"] * 4, texts) for ch, v in out.items(): assert np.all(np.isfinite(v)), f"{ch} produced non-finite values: {v}" assert np.allclose(v, v[0]), f"{ch} should be constant when all gated" assert eng.last_stats.frac_groups_degenerate == 1.0 def test_duplicates_get_low_marginal_within_group(): """Two identical stories should each be worth little in the log-det channel.""" j = FakeJudge() eng = _engine("E2", judge=j) dup = _story("DUP", 250, 7) texts = [dup, dup] + [_story("X", 250, i) for i in range(3, 7)] out = eng.compute(["p"] * 6, texts) m = out["marginal"] assert m[0] < m[2:].mean() and m[1] < m[2:].mean(), \ f"duplicates were not penalized in the marginal channel: {m}" def test_two_groups_are_scored_independently(): j = FakeJudge() eng = _engine("E1", judge=j) texts = [_story("A", 250, i) for i in range(4)] + [_story("B", 250, i + 10) for i in range(4)] prompts = ["p1"] * 4 + ["p2"] * 4 out = eng.compute(prompts, texts) assert len(out["deviation"]) == 8 assert eng.last_stats.n == 8 def test_reward_funcs_match_channel_order(): eng = _engine("E2") fns = eng.make_reward_funcs() assert [f.__name__ for f in fns] == \ ["quality_reward", "deviation_reward", "marginal_reward"] def test_engine_memoizes_within_a_batch(): """TRL calls one reward fn per channel; the judge+embedder must run once.""" j = FakeJudge() eng = _engine("E2", judge=j) texts = [_story("A", 250, i) for i in range(4)] comps = [[{"role": "assistant", "content": t}] for t in texts] prompts = [[{"role": "user", "content": "p"}]] * 4 fns = eng.make_reward_funcs() for f in fns: f(comps, prompts=prompts) assert j.calls == 4, f"judge called {j.calls} times; memoization failed" if __name__ == "__main__": import sys, traceback fns = [(n, f) for n, f in sorted(globals().items()) if n.startswith("test_") and callable(f)] bad = 0 for n, f in fns: try: f(); print(f" PASS {n}") except Exception: bad += 1; print(f" FAIL {n}"); traceback.print_exc() print(f"\n{len(fns)-bad}/{len(fns)} passed") sys.exit(1 if bad else 0)