18. 调试强化学习训练

#@save
import copy
import dataclasses
import torch
import torch.nn.functional as F
import agentrl as arl

torch.set_num_threads(4)                     # the check host is shared (see Chapter 10)
tok = arl.build_tokenizer(arl.make_tasks('all'))
train_tasks, heldout_tasks = arl.make_tasks('train'), arl.make_tasks('heldout')
base = arl.make_rl_base(tok)                 # Chapter 9's starting policy

#@save
def log_metrics(sampled, cfg, update=None):
    """One dashboard row from a step's sampled Trajectory list (+ the update's stats) -> flat dict."""
    n, G = len(sampled), cfg.group_size
    rewards = torch.tensor([t.reward for t in sampled]).view(-1, G)                  # [P, G]
    passed = [bool(t.info['submitted'] and t.info['hidden_pass']) for t in sampled]
    lengths = [float(t.action_mask.sum()) for t in sampled]                           # policy tokens per episode
    def mean_length(want):
        picked = [length for length, p in zip(lengths, passed) if p == want]
        return sum(picked) / max(len(picked), 1)
    row = dict(arl.behavior_stats(sampled), n_episodes=n,
               zero_adv_frac=(rewards.max(-1).values == rewards.min(-1).values).float().mean().item(),
               malformed=sum(t.info['n_invalid'] for t in sampled) / n,
               hardcode=sum('hardcode_tests' in t.info['ops'] for t in sampled) / n,
               submitted=sum(bool(t.info['submitted']) for t in sampled) / n,
               length_pass=mean_length(True), length_fail=mean_length(False))
    return dict(row, **(update or {}))

#@save
BUGS = {
    'mask_shift': '(1) action mask aligned to ids instead of ids[1:]: each policy span moves one token right',
    'stale_old_logp': '(2) old_logp recomputed from the current weights before every minibatch',
    'entropy_sign': '(3) entropy bonus added with the wrong sign: an entropy penalty of ENTROPY_COEF',
    'sign_flip': '(5) advantages negated',
    'accum_local': '(8) minibatch accumulated over 4 length-sorted microbatches, each token_mean / 4',
    'sampler_temp': '(9) sampler at SAMPLER_TEMP and its log-probs used as old_logp; trainer at cfg.temperature',
    'ref_is_policy': '(10) the KL reference is the live policy object',
}
SAMPLER_TEMP, ENTROPY_COEF = 0.7, 0.3


def prepare_batch(model, batch, cfg, bug=None):
    """What the trainer feeds its loss -> (mask, old_logp [B, T-1], advantages [B]), with an optional fault."""
    ids, mask = batch['ids'], batch['action_mask']                         # [B, T], [B, T-1]
    if bug == 'mask_shift':
        mask = F.pad(mask, (1, 0))[:, :-1]
    with torch.no_grad():
        old_logp = arl.selected_log_probs(model(ids)[:, :-1] / cfg.temperature, ids[:, 1:])
    if bug == 'sampler_temp':
        old_logp = batch['sample_logp']
    advantages = arl.compute_advantages(batch, cfg)
    return mask, old_logp, -advantages if bug == 'sign_flip' else advantages


def make_ref(model, cfg, bug=None):
    """The frozen KL reference (None when cfg.kl_coef is 0); the live model itself under 'ref_is_policy'."""
    if not cfg.kl_coef:
        return None
    return model if bug == 'ref_is_policy' else copy.deepcopy(model).eval()

#@save
def minibatch_backward(model, batch, rows, mask, old_logp, advantages, cfg, bug=None, ref_model=None):
    """Loss and backward for one minibatch, no optimizer step -> stats dict (clip_frac, ..., entropy, loss)."""
    ids = batch['ids']
    parts = rows[mask[rows].sum(-1).argsort()].chunk(4) if bug == 'accum_local' else [rows]
    logs = []
    for part in parts:
        logits = model(ids[part])[:, :-1] / cfg.temperature                    # [b, T-1, V]
        logp = arl.selected_log_probs(logits, ids[part, 1:])
        loss, stats = arl.policy_loss(logp, old_logp[part], advantages[part], mask[part], cfg)
        if bug == 'entropy_sign':
            loss = loss + ENTROPY_COEF * arl.entropy_bonus(logits, mask[part])
        if ref_model is not None:
            with torch.no_grad():
                ref_logp = arl.selected_log_probs(ref_model(ids[part])[:, :-1] / cfg.temperature, ids[part, 1:])
            kl = arl.masked_mean(arl.kl_penalty(logp, ref_logp, cfg.kl_kind), mask[part])
            loss, stats['kl_ref'] = loss + cfg.kl_coef * kl, kl.item()
        (loss / len(parts)).backward()
        stats['entropy'] = arl.masked_mean(arl.entropy_from_logits(logits.detach()), mask[part]).item()
        logs.append(dict(stats, loss=loss.item()))
    return {k: sum(s[k] for s in logs) / len(logs) for k in logs[0]}

#@save
def debug_update(model, optimizer, batch, cfg, bug=None, ref_model=None, generator=None):
    """arl.rl_update with diagnostics and an optional fault from BUGS -> dict of mean stats."""
    mask, old_logp, advantages = prepare_batch(model, batch, cfg, bug)
    trainer_logp = prepare_batch(model, batch, cfg)[1] if bug else old_logp      # clean recomputation
    gap = (trainer_logp - batch['sample_logp']).abs()[mask > 0]                   # [N trained tokens]
    logs = []
    for _ in range(cfg.epochs):
        for rows in torch.randperm(len(advantages), generator=generator).chunk(cfg.minibatches):
            if bug == 'stale_old_logp':
                old_logp = prepare_batch(model, batch, cfg)[1]
            optimizer.zero_grad()
            stats = minibatch_backward(model, batch, rows, mask, old_logp, advantages, cfg, bug, ref_model)
            stats['grad_norm'] = float(torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm))
            optimizer.step()
            logs.append(stats)
    out = {k: sum(s[k] for s in logs) / len(logs) for k in logs[0]}
    return dict(out, logp_gap=gap.mean().item(), loss_tokens=mask.sum().item() / len(mask))

#@save
def train_logged(model, tasks, tok, cfg, steps, bug=None, eval_tasks=None, eval_every=10, env_factory=None):
    """arl.train_rl with log_metrics rows and an optional fault from BUGS; trains in place -> history."""
    torch.manual_seed(cfg.seed)
    gen = torch.Generator().manual_seed(cfg.seed)
    env_factory = env_factory or (lambda: arl.make_env(cfg))
    sampler_cfg = dataclasses.replace(cfg, temperature=SAMPLER_TEMP) if bug == 'sampler_temp' else cfg
    ref_model = make_ref(model, cfg, bug)
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
    history = []
    for step in range(steps):
        kept, sampled = arl.sample_step(model, env_factory, tasks, tok, sampler_cfg, gen)
        batch = arl.collate(kept, tok.pad_id)
        batch['n_turns'] = torch.tensor([float(t.info['n_turns']) for t in kept])
        row = log_metrics(sampled, cfg, debug_update(model, optimizer, batch, cfg, bug, ref_model, gen))
        if eval_tasks is not None and (step % eval_every == 0 or step == steps - 1):
            row['eval_pass'] = arl.evaluate(model, eval_tasks, tok, 4, cfg.temperature,
                                            torch.Generator().manual_seed(step), cfg, env_factory)['pass_rate']
        history.append(dict(row, step=step))
    return history

reference, logged = copy.deepcopy(base), copy.deepcopy(base)
hist_ref = arl.train_rl(reference, train_tasks, tok, arl.RLConfig(), steps=5)
hist_log = train_logged(logged, train_tasks, tok, arl.RLConfig(), steps=5)
same_params = all(torch.equal(p, q) for p, q in zip(reference.parameters(), logged.parameters()))
same_stats = all(a[k] == b[k] for a, b in zip(hist_ref, hist_log) for k in a)
print(f'parameters bit-identical: {same_params} | all {len(hist_ref[0])} train_rl statistics identical: {same_stats}')

parameters bit-identical: True | all 17 train_rl statistics identical: True

RUNS = {}
def run(name, cfg=arl.RLConfig(), bug=None, steps=24):
    """Train a fresh copy of the RL base agent with train_logged; keep its history in RUNS[name]."""
    RUNS[name] = train_logged(copy.deepcopy(base), train_tasks, tok, cfg, steps, bug, heldout_tasks, eval_every=8)

def show(names, keys):
    """Per run: each metric as mean(first 4 steps) -> mean(last 4 steps), and the last held-out pass rate."""
    for name in names:
        h = RUNS[name]
        mean = lambda k, part: sum(r[k] for r in part) / len(part)
        cells = ' | '.join(f'{k} {mean(k, h[:4]):.3g} -> {mean(k, h[-4:]):.3g}' for k in keys)
        print(f"{name:<17}| {cells} | held-out {h[-1]['eval_pass']:.2f}")

run('healthy')

TAXONOMY = ['reward', 'eval_pass', 'entropy', 'approx_kl', 'clip_frac', 'ratio_max', 'grad_norm', 'length',
            'zero_adv_frac', 'reads', 'malformed', 'hardcode', 'logp_gap']
healthy = RUNS['healthy']
for key in TAXONOMY:
    values = sorted(r[key] for r in healthy if key in r)
    print(f'{key:<14} step 0 {healthy[0][key]:9.3g} | last {healthy[-1][key]:9.3g} |'
          f' min {values[0]:9.3g} | median {values[len(values) // 2]:9.3g} | max {values[-1]:9.3g}')

reward         step 0     0.172 | last     0.656 | min     0.172 | median     0.438 | max     0.859
eval_pass      step 0       0.1 | last       0.7 | min       0.1 | median     0.525 | max       0.7
entropy        step 0     0.565 | last      0.21 | min     0.169 | median     0.355 | max     0.565
approx_kl      step 0   0.00322 | last  0.000681 | min  9.61e-05 | median  0.000303 | max   0.00322
clip_frac      step 0    0.0139 | last         0 | min         0 | median         0 | max    0.0139
ratio_max      step 0      1.21 | last      1.27 | min      1.03 | median      1.11 | max      1.27
grad_norm      step 0     0.204 | last     0.243 | min     0.155 | median     0.226 | max     0.288
length         step 0      13.2 | last      14.4 | min        11 | median      13.7 | max      14.4
zero_adv_frac  step 0     0.375 | last     0.375 | min         0 | median      0.25 | max     0.625
reads          step 0     0.312 | last      1.08 | min     0.312 | median      1.11 | max      1.61
malformed      step 0    0.0625 | last    0.0469 | min    0.0312 | median    0.0781 | max     0.188
hardcode       step 0    0.0781 | last         0 | min         0 | median    0.0156 | max     0.125
logp_gap       step 0  3.05e-08 | last  6.98e-09 | min  6.98e-09 | median  3.07e-08 | max  5.02e-08

run('mask_shift', bug='mask_shift')
run('stale_old_logp', bug='stale_old_logp')
run('sampler_temp', bug='sampler_temp')
show(['healthy', 'mask_shift', 'stale_old_logp', 'sampler_temp'],
     ['reward', 'logp_gap', 'clip_frac', 'ratio_max', 'grad_norm'])

healthy          | reward 0.191 -> 0.656 | logp_gap 3.79e-08 -> 8.86e-09 | clip_frac 0.00385 -> 0 | ratio_max 1.12 -> 1.14 | grad_norm 0.191 -> 0.229 | held-out 0.70
mask_shift       | reward 0.188 -> 0.52 | logp_gap 3.75 -> 3.1 | clip_frac 0.0129 -> 0.0437 | ratio_max 1.23 -> 2.96 | grad_norm 0.478 -> 1.5 | held-out 0.47
stale_old_logp   | reward 0.207 -> 0.715 | logp_gap 4.32e-08 -> 1.51e-08 | clip_frac 0 -> 0 | ratio_max 1 -> 1 | grad_norm 0.225 -> 0.238 | held-out 0.68
sampler_temp     | reward 0.227 -> 0.156 | logp_gap 0.0752 -> 0.0285 | clip_frac 0.0218 -> 0 | ratio_max 2.65 -> 2.96 | grad_norm 0.202 -> 0 | held-out 0.15

show(['healthy', 'mask_shift', 'sampler_temp'], ['entropy', 'malformed', 'reads', 'zero_adv_frac'])

healthy          | entropy 0.531 -> 0.195 | malformed 0.0625 -> 0.0664 | reads 0.473 -> 1.09 | zero_adv_frac 0.312 -> 0.375 | held-out 0.70
mask_shift       | entropy 0.579 -> 0.795 | malformed 0.043 -> 0.121 | reads 0.418 -> 1.2 | zero_adv_frac 0.312 -> 0.0938 | held-out 0.47
sampler_temp     | entropy 0.47 -> 0.156 | malformed 0 -> 0 | reads 0.0859 -> 0.00391 | zero_adv_frac 0.281 -> 1 | held-out 0.15

check_cfg = arl.RLConfig(kl_coef=0.1)          # with a KL term, so that Section 18.6's reference checks apply
def probe(bug=None, seed=1):
    """64 episodes of the base agent on 8 training tasks, sampled as the (possibly faulty) pipeline samples."""
    cfg = dataclasses.replace(check_cfg, temperature=SAMPLER_TEMP) if bug == 'sampler_temp' else check_cfg
    return arl.sample_groups(base, lambda: arl.make_env(cfg), train_tasks[:8], tok, cfg,
                             torch.Generator().manual_seed(seed))

batch = arl.collate(probe('sampler_temp'), tok.pad_id)                       # sampled at T = 0.7
mask, old_logp, adv = prepare_batch(base, batch, check_cfg, 'sampler_temp')
ratio = (prepare_batch(base, batch, check_cfg)[1] - old_logp).exp()          # pi_T=1 / mu_T=0.7, [B, T-1]
rare = old_logp.exp() < 0.5                                                  # the sampler gave them < 0.5
positive = (adv > 0)[:, None] & (mask > 0)                                   # tokens the update should reinforce
for name, sel in [('rare (mu < 0.5)', rare & positive), ('common (mu >= 0.5)', ~rare & positive)]:
    print(f'{name:<19} tokens with A > 0: {int(sel.sum()):4d} | mean ratio {ratio[sel].mean():.2f} |'
          f' clipped (ratio > 1.2) {(ratio[sel] > 1.2).float().mean():.2f}')

rare (mu < 0.5)     tokens with A > 0:   17 | mean ratio 1.17 | clipped (ratio > 1.2) 0.12
common (mu >= 0.5)  tokens with A > 0:   91 | mean ratio 0.95 | clipped (ratio > 1.2) 0.00

run('sign_flip', bug='sign_flip', steps=8)
run('lr 1e-2', arl.RLConfig(lr=1e-2), steps=8)
run('accum_local', bug='accum_local')
show(['healthy', 'sign_flip', 'lr 1e-2', 'accum_local'], ['reward', 'approx_kl', 'entropy', 'malformed', 'turns'])
for name in ['healthy', 'lr 1e-2']:
    print(f'{name:<17}', ' '.join(f"{r['approx_kl']:.4f}" for r in RUNS[name][:8]), '| approx_kl, steps 0-7')

healthy          | reward 0.191 -> 0.656 | approx_kl 0.00112 -> 0.000442 | entropy 0.531 -> 0.195 | malformed 0.0625 -> 0.0664 | turns 3.52 -> 4.26 | held-out 0.70
sign_flip        | reward 0.125 -> 0.0273 | approx_kl 0.00145 -> 0.000167 | entropy 0.551 -> 0.501 | malformed 0.0664 -> 0.0312 | turns 4.21 -> 5.48 | held-out 0.03
lr 1e-2          | reward 0.043 -> 0 | approx_kl 0.0953 -> 0.00214 | entropy 0.417 -> 1.08 | malformed 0.762 -> 2.89 | turns 4.75 -> 6 | held-out 0.00
accum_local      | reward 0.219 -> 0.668 | approx_kl 0.00129 -> 0.000197 | entropy 0.499 -> 0.204 | malformed 0.0391 -> 0.16 | turns 3.64 -> 4.38 | held-out 0.65
healthy           0.0032 0.0007 0.0003 0.0003 0.0006 0.0005 0.0006 0.0002 | approx_kl, steps 0-7
lr 1e-2           0.3098 0.0642 0.0049 0.0024 0.0038 0.0021 0.0018 0.0009 | approx_kl, steps 0-7

run('entropy_sign', bug='entropy_sign')
show(['healthy', 'entropy_sign'], ['reward', 'entropy', 'reads', 'zero_adv_frac'])

healthy          | reward 0.191 -> 0.656 | entropy 0.531 -> 0.195 | reads 0.473 -> 1.09 | zero_adv_frac 0.312 -> 0.375 | held-out 0.70
entropy_sign     | reward 0.176 -> 0.176 | entropy 0.445 -> 0.0268 | reads 0.191 -> 0.00781 | zero_adv_frac 0.344 -> 0.844 | held-out 0.20

run('healthy, 12 turns', arl.RLConfig(max_turns=12), steps=16)
run('length bias', arl.RLConfig(max_turns=12, loss_mode='seq_mean_token_mean'), steps=16)
show(['healthy, 12 turns', 'length bias'], ['reward', 'length_pass', 'length_fail', 'turns', 'reads'])

healthy, 12 turns| reward 0.23 -> 0.488 | length_pass 13.9 -> 13.2 | length_fail 16.1 -> 12.8 | turns 4.34 -> 3.98 | reads 0.57 -> 1.08 | held-out 0.55
length bias      | reward 0.195 -> 0.512 | length_pass 14.3 -> 12.7 | length_fail 20.8 -> 33.9 | turns 5.45 -> 6.77 | reads 0.75 -> 2.71 | held-out 0.55

run('public reward', arl.RLConfig(reward='public'))
show(['healthy', 'public reward'], ['reward', 'pass_rate', 'hardcode', 'reads'])

healthy          | reward 0.191 -> 0.656 | pass_rate 0.191 -> 0.656 | hardcode 0.0703 -> 0 | reads 0.473 -> 1.09 | held-out 0.70
public reward    | reward 0.273 -> 0.969 | pass_rate 0.176 -> 0.00781 | hardcode 0.125 -> 0.969 | reads 0.418 -> 0.0352 | held-out 0.00

run('KL 0.1', arl.RLConfig(kl_coef=0.1), steps=8)
run('ref_is_policy', arl.RLConfig(kl_coef=0.1), bug='ref_is_policy', steps=8)
show(['KL 0.1', 'ref_is_policy'], ['reward', 'kl_ref'])
print('ref_is_policy identical to the run without KL:',
      all(a['reward'] == b['reward'] and a['grad_norm'] == b['grad_norm'] for a, b in zip(RUNS['ref_is_policy'], healthy)))

KL 0.1           | reward 0.203 -> 0.27 | kl_ref 0.0142 -> 0.0571 | held-out 0.23
ref_is_policy    | reward 0.191 -> 0.277 | kl_ref 0 -> 0 | held-out 0.28
ref_is_policy identical to the run without KL: True

#@save
def pad_batch(batch, extra, pad_id=0):
    """The same batch with `extra` columns of right padding on every [B, T] and [B, T-1] tensor."""
    return dict(batch, **{k: F.pad(batch[k], (0, extra), value=pad_id if k == 'ids' else 0)
                          for k in ['ids', 'attention_mask', 'action_mask', 'sample_logp']})


def flat_grad(model, batch, mask, old_logp, advantages, cfg, bug=None):
    """The trainer's backward on the whole batch, no optimizer step -> (gradient vector, stats)."""
    model.zero_grad()
    stats = minibatch_backward(model, batch, torch.arange(len(advantages)), mask, old_logp, advantages, cfg, bug)
    grad = torch.cat([p.grad.flatten() for p in model.parameters()])
    model.zero_grad()
    return grad, stats


def step_checks(model, batch, owned, trainer_logp, cfg, bug=None):
    """Invariants of one real update on a copy of model -> dict of (measured value, passed)."""
    trained = copy.deepcopy(model)
    ref = make_ref(trained, cfg, bug)
    stats = debug_update(trained, torch.optim.Adam(trained.parameters(), lr=cfg.lr), batch, cfg, bug, ref,
                         torch.Generator().manual_seed(0))
    new_logp = prepare_batch(trained, batch, cfg)[1]
    moved = ((new_logp - trainer_logp) * owned).sum(-1)                         # change of log pi(episode), [B]
    corr = torch.corrcoef(torch.stack([moved, arl.compute_advantages(batch, cfg)]))[0, 1].item()
    kl = 0.0 if ref is None else arl.masked_mean(arl.kl_penalty(new_logp, prepare_batch(ref, batch, cfg)[1]),
                                                 owned).item()
    return {'old_logp frozen in the update': (stats['approx_kl'], stats['approx_kl'] > 0),
            'update raises log pi where A > 0': (corr, corr > 0),
            'KL to reference > 0 after a step': (kl, ref is None or kl > 0)}

#@save
def sanity_checks(model, trajs, tok, cfg, bug=None):
    """Invariants of the instrumented trainer on one batch at step 0 -> dict check -> (measured value, passed)."""
    batch = arl.collate(trajs, tok.pad_id)
    batch['n_turns'] = torch.tensor([float(t.info['n_turns']) for t in trajs])
    mask, old_logp, adv = prepare_batch(model, batch, cfg, bug)
    owned = (arl.pad_rows([t.owner[1:] for t in trajs], mask.shape[1], -1) == arl.OWNER['policy']).float()
    trainer_logp = prepare_batch(model, batch, cfg)[1]
    gap = (trainer_logp - batch['sample_logp']).abs()[mask > 0].max().item()
    full, _ = flat_grad(model, batch, mask, old_logp, adv, cfg, None if bug == 'accum_local' else bug)
    grad, stats = flat_grad(model, batch, mask, old_logp, adv, cfg, bug)          # the trainer's own path
    padded = pad_batch(batch, 8, tok.pad_id)
    _, pad_stats = flat_grad(model, padded, *prepare_batch(model, padded, cfg, bug), cfg, bug)
    ratio_err = max(stats['ratio_max'] - 1, 1 - stats['ratio_min'])
    group_mean = adv.view(-1, cfg.group_size).mean(-1).abs().max().item()
    pad_err = abs(pad_stats['loss'] - stats['loss'])
    accum_err = ((grad - full).norm() / full.norm()).item()
    checks = {'mask = policy tokens': (int((mask != owned).sum()), bool((mask == owned).all())),
              'sampler = trainer log-prob': (gap, gap < 1e-4),
              'ratio = 1 before any step': (ratio_err, ratio_err < 1e-4),
              'group advantage mean = 0': (group_mean, group_mean < 1e-5),
              'loss invariant to padding': (pad_err, pad_err < 1e-5),
              'accumulated = one-pass gradient': (accum_err, accum_err < 1e-4)}
    return dict(checks, **step_checks(model, batch, owned, trainer_logp, cfg, bug))

for name, (value, ok) in sanity_checks(base, probe(), tok, check_cfg).items():
    print(f'{name:<34} {value:10.3g}  {"ok" if ok else "FAIL"}')

mask = policy tokens                        0  ok
sampler = trainer log-prob           1.31e-06  ok
ratio = 1 before any step                   0  ok
group advantage mean = 0             2.24e-08  ok
loss invariant to padding            7.45e-09  ok
accumulated = one-pass gradient             0  ok
old_logp frozen in the update         0.00223  ok
update raises log pi where A > 0        0.514  ok
KL to reference > 0 after a step       0.0124  ok

MATRIX = {bug: sanity_checks(base, probe(bug), tok, check_cfg, bug) for bug in [None, *BUGS]}
names = list(MATRIX[None])
print(f'{"":<15}', ' '.join(f'{i + 1:>4}' for i in range(len(names))))
for bug, checks in MATRIX.items():
    print(f'{str(bug):<15}', ' '.join(f'{"ok" if ok else "FAIL":>4}' for _, ok in checks.values()))

                   1    2    3    4    5    6    7    8    9
None              ok   ok   ok   ok   ok   ok   ok   ok   ok
mask_shift      FAIL FAIL   ok   ok FAIL   ok   ok   ok   ok
stale_old_logp    ok   ok   ok   ok   ok   ok FAIL   ok   ok
entropy_sign      ok   ok   ok   ok   ok   ok   ok   ok   ok
sign_flip         ok   ok   ok   ok   ok   ok   ok FAIL   ok
accum_local       ok   ok   ok   ok   ok FAIL   ok   ok   ok
sampler_temp      ok FAIL FAIL   ok   ok   ok   ok   ok   ok
ref_is_policy     ok   ok   ok   ok   ok   ok   ok   ok FAIL

for name in ['healthy', 'mask_shift']:
    h = RUNS[name]
    print(f'{name:<11}| reward', ' '.join(f"{h[i]['reward']:.2f}" for i in range(0, 24, 4)),
          f"| held-out {h[-1]['eval_pass']:.2f} | logp_gap {h[0]['logp_gap']:.2g} | grad_norm {h[0]['grad_norm']:.2f}")

healthy    | reward 0.17 0.34 0.36 0.44 0.53 0.72 | held-out 0.70 | logp_gap 3e-08 | grad_norm 0.20
mask_shift | reward 0.17 0.27 0.39 0.41 0.39 0.50 | held-out 0.47 | logp_gap 3.5 | grad_norm 0.48

h = RUNS['public reward']
for step in [0, 8, 16, 23]:
    r = h[step]
    print(f"step {step:2d} | reward {r['reward']:.3f} | entropy {r['entropy']:.3f} | grad_norm {r['grad_norm']:.3f}"
          f" | approx_kl {r['approx_kl']:.1e} | held-out hidden pass {r['eval_pass']:.3f}")

step  0 | reward 0.234 | entropy 0.565 | grad_norm 0.252 | approx_kl 3.3e-03 | held-out hidden pass 0.175
step  8 | reward 0.750 | entropy 0.310 | grad_norm 0.404 | approx_kl 4.7e-04 | held-out hidden pass 0.125
step 16 | reward 0.969 | entropy 0.115 | grad_norm 0.155 | approx_kl 1.4e-04 | held-out hidden pass 0.025
step 23 | reward 0.969 | entropy 0.092 | grad_norm 0.200 | approx_kl 2.2e-05 | held-out hidden pass 0.000

def signature(h):
    """What a rerun must match: per-step reward, length, and gradient norm."""
    return [(r['reward'], r['length'], r['grad_norm']) for r in h]

again = train_logged(copy.deepcopy(base), train_tasks, tok, arl.RLConfig(), 6)
nudged = copy.deepcopy(base)
with torch.no_grad():                                  # relative 1e-5 noise on every weight
    for p in nudged.parameters():
        p.mul_(1 + 1e-5 * torch.randn(p.shape, generator=torch.Generator().manual_seed(0)))
perturbed = train_logged(nudged, train_tasks, tok, arl.RLConfig(), 6)
other_seed = train_logged(copy.deepcopy(base), train_tasks, tok, arl.RLConfig(seed=1), 6)
print('same seed, 6 steps, identical:', signature(again) == signature(healthy[:6]))
for name, h in [('weights x (1 + 1e-5 noise)', perturbed), ('seed 1', other_seed)]:
    print(f'{name:<26} rewards equal at {sum(a["reward"] == b["reward"] for a, b in zip(h, again))}/6 steps |'
          f' max |grad_norm diff| {max(abs(a["grad_norm"] - b["grad_norm"]) for a, b in zip(h, again)):.1e}')

same seed, 6 steps, identical: True
weights x (1 + 1e-5 noise) rewards equal at 6/6 steps | max |grad_norm diff| 8.3e-06
seed 1                     rewards equal at 0/6 steps | max |grad_norm diff| 1.7e-01

COMMITS = ['cache eval tasks', 'log entropy per turn', 'vectorize collate', 'inline advantage helper',
           'bf16 autocast in sampler', 'rename config fields', 'new task split', 'faster sampler']
BAD = 3                                                # simulated history: commit 3 introduced a sign error
def checks_pass(commit):
    """Does the code at this commit pass every invariant? (The search below never reads BAD.)"""
    bug = 'sign_flip' if commit >= BAD else None
    return all(ok for _, ok in sanity_checks(base, probe(), tok, check_cfg, bug).values())

good, bad = 0, len(COMMITS) - 1                         # commit 0 known good, the last one known bad
while bad - good > 1:
    mid = (good + bad) // 2
    passed = checks_pass(mid)
    print(f'test commit {mid} ({COMMITS[mid]}): {"pass" if passed else "fail"}')
    good, bad = (mid, bad) if passed else (good, mid)
print(f'first bad commit: {bad} ({COMMITS[bad]})')

test commit 3 (inline advantage helper): fail
test commit 1 (log entropy per turn): pass
test commit 2 (vectorize collate): pass
first bad commit: 3 (inline advantage helper)