code/notebooks/_m1_src/0107-cispo.py · viai957/align at main

align / code /notebooks /_m1_src /0107-cispo.py
Add files using upload-large-folder tool
63f41a6 verified 16 days ago
22.7 kB
# %% [markdown]
# # 0107 - CISPO (MiniMax-M1 Sec. 3.1): clip the importance weight, not the update
# Companion to `notes/0107 - CISPO.pdf`.  fp64; exact enumeration on the 121-sequence toy LM.
# %%
import sys, math
sys.path.insert(0, "..")
from _m1_lib import *
import numpy as np
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
R = default_reward(lm)
th_old = lm.theta0.clone(); q = lm.probs(th_old); N, L = lm.N, lm.L
b = float(q @ R)
A_tok = (R - b)[:, None] * lm.mask                       # outcome advantage, same for all tokens of a sequence
lq = lm.token_logprobs(th_old)
Ttens = token_score_tensor(lm, th_old)                    # score of each token at THETA_OLD (used to build gradients analytically)
eh, el = 0.28, 0.2
# %% [markdown]
# ## 1. Gradient weights.  Every method has gradient  sum_t  w(r_t, A_t) * A_t * grad log pi_theta(o_t | s_t)
#   PPO/GRPO (clip lo/hi):   w = r * 1[live]       live = (A>0 & r<1+eh) | (A<0 & r>1-el)
#   CISPO:                   w = min(r, 1+eh)      (A-independent, always live)      [loss = sg(min(r,1+eh)) * A * log pi]
#   token IS-PG (no clip):   w = r
# We verify against autograd of the actual loss functions on the toy model at a perturbed theta.
# %%
def ratio_and_lp(th):
    lp = lm.token_logprobs(th)
    r = (lp - lq).exp() * lm.mask + (1 - lm.mask)
    return r, lp
def loss_ppo(th):
    r, lp = ratio_and_lp(th)
    f = torch.where(A_tok > 0, A_tok * torch.clamp(r, max=1 + eh), A_tok * torch.clamp(r, min=1 - el))
    return -(q[:, None] * f * lm.mask).sum()
def loss_cispo(th, eps_hi=eh):
    r, lp = ratio_and_lp(th)
    rhat = torch.clamp(r, max=1 + eps_hi).detach()                    # sg( clip(r) )
    return -(q[:, None] * rhat * A_tok * lp * lm.mask).sum()
def loss_isPG(th):
    r, lp = ratio_and_lp(th)
    return -(q[:, None] * r.detach() * A_tok * lp * lm.mask).sum()
def weights(th, method):
    r, _ = ratio_and_lp(th); r = r.detach()
    if method == 'ppo': live = ((A_tok > 0) & (r < 1 + eh)) | ((A_tok < 0) & (r > 1 - el)); return r * live
    if method == 'cispo': return torch.clamp(r, max=1 + eh)
    if method == 'is': return r
def grad_from_weights(th, w):
    # sum_{o,t} q(o) w A  grad log pi_theta(o_t|s_t)  with the score evaluated at theta (autograd through lp)
    t = th.clone().requires_grad_(True)
    lp = lm.token_logprobs(t)
    return torch.autograd.grad((q[:, None] * w * A_tok * lp * lm.mask).sum(), t)[0]
th = perturb(th_old, 1.5, seed=11)
r_now, _ = ratio_and_lp(th)
print("ratio range over (o,t):  min %.3f  max %.3f" % (float(r_now[lm.mask > 0].min()), float(r_now[lm.mask > 0].max())))
for name, lossf in (("ppo", loss_ppo), ("cispo", loss_cispo), ("is", loss_isPG)):
    t = th.clone().requires_grad_(True)
    g_auto = torch.autograd.grad(lossf(t), t)[0]
    g_w = -grad_from_weights(th, weights(th, name))
    check(f"{name}: autograd of the loss == sum_t w(r,A) A grad log pi  (weights table)", maxabs(g_auto, g_w), 1e-13)
# %% [markdown]
# ## 2. Exactly which tokens does PPO/GRPO silence and CISPO keep?
# %%
live = ((A_tok > 0) & (r_now < 1 + eh)) | ((A_tok < 0) & (r_now > 1 - el))
sel = lm.mask > 0
pw = q[:, None] * lm.mask                                  # probability weight of each (o,t) token occurrence
tot = float(pw.sum())
zeroed = sel & ~live
cells = {"A>0, r>1+eh  (silenced by PPO)": sel & (A_tok > 0) & (r_now > 1 + eh),
         "A<0, r<1-el  (silenced by PPO)": sel & (A_tok < 0) & (r_now < 1 - el),
         "A<0, r>1+eh  (PPO weight r unbounded; CISPO caps at 1+eh)": sel & (A_tok < 0) & (r_now > 1 + eh)}
for k, m in cells.items(): print(f"   {k:60s}: {int(m.sum()):4d} token slots, probability mass {float((pw*m).sum())/tot:.4f}")
Wc = weights(th, 'cispo'); Wp = weights(th, 'ppo')
silenced_by_ppo_kept_by_cispo = sel & (Wp == 0) & (Wc > 0)
check("{tokens with PPO weight 0 and CISPO weight>0} == cells 1 U 2 exactly (symmetric difference size)", float((silenced_by_ppo_kept_by_cispo ^ (cells["A>0, r>1+eh  (silenced by PPO)"] | cells["A<0, r<1-el  (silenced by PPO)"])).sum()), 0.0)
check("CISPO never silences a token (min weight over all (o,t) with A != 0)", abs(float(Wc[sel].min()) - float(r_now[sel].min())) if float(r_now[sel].min()) < 1 + eh else 0.0, 1e-15)
check("PPO silenced set == the MiniMax-M1 mask M_{i,t}=0 of Eq. 7 (symmetric difference size)",
      float((zeroed ^ (((A_tok > 0) & (r_now > 1 + eh)) | ((A_tok < 0) & (r_now < 1 - el))) & sel).sum()), 0.0)
# weight table
print("\n   gradient weight w(r,A)  (loss gradient = w * A * grad log pi),  eps_hi=0.28, eps_lo=0.2")
print("      r      PPO(A>0)  PPO(A<0)   CISPO   unclipped-IS")
for r in (0.3, 0.8, 1.0, 1.28, 2.0, 5.0, 20.0):
    wp_pos = r if r < 1 + eh else 0.0; wp_neg = r if r > 1 - el else 0.0
    print(f"    {r:5.2f}   {wp_pos:7.2f}   {wp_neg:7.2f}   {min(r, 1+eh):6.2f}   {r:7.2f}")
# %% [markdown]
# ## 3. Unified mask formulation (paper Eq. 6-7):  J = sg(r_hat) A log pi * M ;  with r_hat = r it reproduces PPO's gradient
# %%
def loss_unify(th, rhat_clip=None):
    r, lp = ratio_and_lp(th)
    M = (~(((A_tok > 0) & (r > 1 + eh)) | ((A_tok < 0) & (r < 1 - el)))).double()
    rhat = r if rhat_clip is None else torch.clamp(r, max=rhat_clip)
    return -(q[:, None] * rhat.detach() * A_tok * lp * M * lm.mask).sum()
t = th.clone().requires_grad_(True); g_u = torch.autograd.grad(loss_unify(t), t)[0]
t = th.clone().requires_grad_(True); g_p = torch.autograd.grad(loss_ppo(t), t)[0]
check("unified formulation with r_hat=r and PPO mask == PPO gradient", maxabs(g_u, g_p), 1e-13)
# without clipping CISPO reduces to the standard (token-level IS) policy gradient; at theta_old it is the on-policy gradient
t = th.clone().requires_grad_(True); g_c_inf = torch.autograd.grad(loss_cispo(t, eps_hi=1e9), t)[0]
t = th.clone().requires_grad_(True); g_is = torch.autograd.grad(loss_isPG(t), t)[0]
check("CISPO with eps_hi=inf == token-level IS policy gradient", maxabs(g_c_inf, g_is), 1e-14)
t = th_old.clone().requires_grad_(True); g_c0 = torch.autograd.grad(loss_cispo(t), t)[0]
t2 = th_old.clone().requires_grad_(True); gJ = torch.autograd.grad((lm.probs(t2) * R).sum(), t2)[0]
check("at theta_old CISPO == exact policy gradient grad J (r=1, nothing clipped)", maxabs(-g_c0, gJ), 1e-14)
# %% [markdown]
# ## 4. Bias of CISPO relative to the (token-level) IS gradient, and the bound  ||bias|| <= sum_t E[ (r - 1 - eps)_+ |A| ||score|| ]
# %%
def cispo_bias(th, eps):
    r, _ = ratio_and_lp(th); r = r.detach()
    Tth = token_score_tensor(lm, th)                                   # score at theta
    diff = (torch.clamp(r, max=1 + eps) - r)                            # r_hat - r  (<= 0)
    return (q[:, None, None] * (diff * A_tok)[:, :, None] * Tth).sum((0, 1)), Tth
for eps_ in (0.05, 0.1, 0.28):
    bias, Tth = cispo_bias(th, eps_)
    r, _ = ratio_and_lp(th)
    bound = float((q[:, None] * (r - 1 - eps_).clamp(min=0) * A_tok.abs() * Tth.norm(dim=2) * lm.mask).sum())
    print(f"   eps_hi={eps_:4.2f}:  ||bias|| = {float(bias.norm()):.4e}   bound = {bound:.4e}   ||token-IS gradient|| = {float((q[:,None,None]*(r*A_tok)[:,:,None]*Tth).sum((0,1)).norm()):.4e}")
    assert float(bias.norm()) <= bound + 1e-15
check("bias bound holds (3 values of eps_hi)", 0.0, 0.0)
# for comparison: how far each method's gradient is from the token-level IS gradient
t_ = th.clone().requires_grad_(True); gI = torch.autograd.grad(loss_isPG(t_), t_)[0]
t_ = th.clone().requires_grad_(True); gP = torch.autograd.grad(loss_ppo(t_), t_)[0]
t_ = th.clone().requires_grad_(True); gC = torch.autograd.grad(loss_cispo(t_), t_)[0]
print(f"   ||g_PPO - g_IS|| = {float((gP-gI).norm()):.4e}   ||g_CISPO - g_IS|| = {float((gC-gI).norm()):.4e}   (eps_hi=0.28, eps_lo=0.2; ||g_IS|| = {float(gI.norm()):.4e})")
# distribution-free per-state bound  E_old[(r-1-eps)_+] = sum_a (pi - (1+eps) pi_old)_+  <=  TV(pi, pi_old)
worst = 0.0
for s in lm.prefixes():
    po, pn = lm.cond_probs(th_old, s), lm.cond_probs(th, s)
    lhs = float(((pn - 1.28 * po).clamp(min=0)).sum()); tv = 0.5 * float((pn - po).abs().sum())
    worst = max(worst, lhs - tv)
check("per-state E_old[(r-1-eps)_+] <= TV(pi,pi_old)  (max violation, must be <= 0)", max(0.0, worst), 0.0)
# variance: gradient weights are bounded for CISPO, unbounded for PPO's live A<0 tokens
r, _ = ratio_and_lp(th)
print(f"   max gradient weight over (o,t):  PPO {float(weights(th,'ppo').max()):.3f}   CISPO {float(weights(th,'cispo').max()):.3f}   unclipped IS {float(weights(th,'is').max()):.3f}")
assert float(weights(th, 'cispo').max()) <= 1 + eh + 1e-15
check("CISPO weights bounded by 1+eps_hi", max(0.0, float(weights(th, 'cispo').max()) - (1 + eh)), 0.0)
# %% [markdown]
# ## 5. Stop-gradient semantics: three look-alike implementations, three different gradients
#   (a) correct:    loss = - sg(min(r,1+eh)) * A * log pi         -> grad = -A*min(r,1+eh) * grad log pi
#   (b) forgot sg:  loss = - min(r,1+eh) * A * log pi             -> on the UNclipped set grad = -A * r * (1 + log pi) * grad log pi : an extra  A r log pi  term
#   (c) ratio form: loss = - min(r,1+eh) * A  (PPO-like)           -> grad = -A * r * 1[r<1+eh] * grad log pi   (zero beyond the clip!)
# %%
lp_old = torch.tensor(-2.3)
def three(lp_new, A, eps=0.28):
    lp = torch.tensor(lp_new, requires_grad=True); r = (lp - lp_old).exp()
    ga = torch.autograd.grad(-(torch.clamp(r, max=1 + eps).detach() * A * lp), lp, retain_graph=True)[0]
    gb = torch.autograd.grad(-(torch.clamp(r, max=1 + eps) * A * lp), lp, retain_graph=True)[0]
    gc = torch.autograd.grad(-(torch.clamp(r, max=1 + eps) * A), lp)[0]
    return float(ga), float(gb), float(gc), float(r.detach())
print("   lp_old=-2.3;   gradient wrt log pi_theta")
print("   lp_new     r     A     (a) CISPO     (b) no sg     (c) ratio-clip")
for lpn, A in ((-2.3, 1.0), (-1.5, 1.0), (-1.5, -1.0), (-3.0, 1.0), (-0.5, 1.0), (-0.5, -1.0)):
    ga, gb, gc, rr = three(lpn, A)
    print(f"   {lpn:6.2f}  {rr:6.2f}  {A:+.0f}   {ga:+9.3f}   {gb:+9.3f}   {gc:+9.3f}")
    expect_a = -A * min(rr, 1.28)
    assert abs(ga - expect_a) < 1e-12
    expect_c = -A * rr * (rr < 1.28)
    assert abs(gc - expect_c) < 1e-12
    expect_b = -A * rr * (1 + lpn) if rr < 1.28 else -A * 1.28           # clamp kills the extra term beyond the clip
    assert abs(gb - expect_b) < 1e-12
check("(a) = -A min(r,1+eh);  (b) = -A r (1+log pi) on the live set;  (c) = -A r 1[r<1+eh]   (6 test points, exact)", 0.0, 0.0)
ga, gb, gc, rr = three(-2.3, 1.0)
print(f"   at r=1 (inside the clip), A=+1, log pi=-2.3: (a)={ga:+.3f}   (b)={gb:+.3f}  <- sign flipped by the spurious log-pi factor (1+log pi = -1.3)")
assert ga < 0 < gb
check("without sg a live token's gradient weight becomes r(1+log pi): wrong sign when pi < 1/e", 0.0, 0.0)
ga, gb, gc, rr = three(-1.5, 1.0)
print(f"   at r={rr:.2f} > 1+eps: (a)={ga:+.3f} capped weight kept;  (b)={gb:+.3f} identical (clamp zeroes the extra term);  (c)={gc:+.3f} silenced")

# %% [markdown]
# ## 6. The 'fork token' mechanism over K = 16 off-policy updates on one batch (MiniMax-M1 setting)
# One decision with three tokens: two common tokens and a RARE token w (pi_old = 0.01) with positive advantage.
# Exact expectation over the old distribution; plain SGD on the logits; K = 16 inner steps.
# %%
def fork(method, K=16, lr=30.0, eh_=0.28, el_=0.2, p_w=0.01):
    pi_old = torch.tensor([0.60, 0.39 - 0.0, p_w]); pi_old = pi_old / pi_old.sum()
    reward = torch.tensor([0.0, 0.0, 1.0]); A = reward - float(pi_old @ reward)       # A_w > 0, A_common < 0
    z = pi_old.log().clone().requires_grad_(True)
    traj, gw = [], []
    for k in range(K):
        pi = torch.softmax(z, 0); r = (pi / pi_old)
        if method == 'ppo':
            f = torch.where(A > 0, A * torch.clamp(r, max=1 + eh_), A * torch.clamp(r, min=1 - el_))
            obj = (pi_old * f).sum()
        else:
            rhat = torch.clamp(r, max=1 + eh_).detach()
            obj = (pi_old * rhat * A * pi.log()).sum()
        g = torch.autograd.grad(obj, z)[0]
        # direct contribution of token w's OWN term to the gradient on its own logit (diagnostic)
        rd = r.detach(); wt = float((1 + eh_) if (method == 'cispo' and rd[2] > 1 + eh_) else (rd[2] if (method == 'cispo' or rd[2] < 1 + eh_) else 0.0))
        gw.append(wt)
        with torch.no_grad(): z += lr * g
        traj.append(float(torch.softmax(z.detach(), 0)[2]))
    return traj, gw, float(r[2].detach())
tp, gwp, rp = fork('ppo'); tc, gwc, rc = fork('cispo')
print("   after update #" + " ".join(f"{k+1:6d}" for k in range(0, 16, 3)))
print("   p_w PPO:   " + " ".join(f"{tp[k]:6.3f}" for k in range(0, 16, 3)))
print("   p_w CISPO: " + " ".join(f"{tc[k]:6.3f}" for k in range(0, 16, 3)))
print("   weight on w's own term PPO:   " + " ".join(f"{gwp[k]:6.2f}" for k in range(0, 16, 3)))
print("   weight on w's own term CISPO: " + " ".join(f"{gwc[k]:6.2f}" for k in range(0, 16, 3)))
n_zero_ppo = sum(1 for x in gwp if x == 0.0)
print(f"   PPO: token w gets zero direct gradient in {n_zero_ppo} of 16 inner steps;  CISPO: {sum(1 for x in gwc if x==0.0)} of 16")
assert n_zero_ppo >= 8 and all(x > 0 for x in gwc)
check("PPO silences the rare token for >= 8 of 16 inner steps, CISPO for none", 0.0, 0.0)
print(f"   final p_w after 16 updates: PPO {tp[-1]:.3f}   CISPO {tc[-1]:.3f}")
# %% [markdown]
# ## 7. Naive vs efficient implementation of the CISPO loss, token-level normalisation by total tokens
# %%
Gs = 8
gen = torch.Generator().manual_seed(5)
idx = torch.multinomial(q, Gs, replacement=True, generator=gen)
Rg = R[idx]; Ag = Rg - Rg.mean()
t = perturb(th_old, 0.8, seed=2).clone().requires_grad_(True)
lp = lm.token_logprobs(t)[idx]; lq_ = lq[idx]; mk = lm.mask[idx]
def cispo_naive(lp, lq_):
    tot, n = 0.0, 0
    for i in range(Gs):
        for tt in range(int(mk[i].sum())):
            r = math.exp(float((lp[i, tt] - lq_[i, tt]).detach()))
            tot = tot + min(r, 1 + eh) * Ag[i] * lp[i, tt]; n += 1
    return -tot / n
def cispo_vec(lp, lq_):
    r = (lp - lq_).exp().detach()
    return -((torch.clamp(r, max=1 + eh) * Ag[:, None] * lp * mk).sum() / mk.sum())
gn = torch.autograd.grad(cispo_naive(lp, lq_), t, retain_graph=True)[0]
gv = torch.autograd.grad(cispo_vec(lp, lq_), t, retain_graph=True)[0]
check("naive per-token loop == vectorised masked CISPO loss (gradient)", maxabs(gn, gv), 1e-14)
# %% [markdown]
# ## 8. Leanstral variant (Mistral technical report, Sec. 2.3): sampler ratio, masking of tool/prompt tokens, outcome broadcast
# Loss there:  `L = -(1/sum_i|o_i|) sum_i sum_t sg(min(pi_theta/pi_sample, c)) * A_i * log pi_theta(a_it | s_it, T_i)`, `T_i` = tool calls and responses used as context only,
# one verifier outcome for every model token.  Toy: two-turn tabular agent (4 actions per turn), turn-2 context = turn-1 action, feedback token = a fixed function of turn 1 (not scored).
# `R ~ Bern(T[y1, y2])`.  Everything below is enumerated exactly over the 32 episodes (y1, y2, R).
# %%
import numpy as _np
_rng = _np.random.default_rng(123)
_A = 4
_T = _rng.uniform(0.05, 0.95, (_A, _A))
def _sm(x): e = _np.exp(x - x.max(-1, keepdims=True)); return e / e.sum(-1, keepdims=True)
def _unpack(th): return th[:_A], th[_A:].reshape(_A, _A)
th_cur = _rng.normal(0, 0.8, _A + _A * _A)
th_smp = th_cur + _rng.normal(0, 0.35, _A + _A * _A)              # sampler = lagged + numerically perturbed copy
def _probs(th): a, b = _unpack(th); return _sm(a), _sm(b)
def _J(th): p1, p2 = _probs(th); return float((p1[:, None] * p2 * _T).sum())
def _score(th, y1, y2):                                           # (s1, s2): grads of log pi(y1), log pi(y2|y1) wrt the 20 logits
    p1, p2 = _probs(th); s1 = _np.zeros(_A + _A * _A); s2 = _np.zeros(_A + _A * _A)
    s1[:_A] = _np.eye(_A)[y1] - p1
    blk = _np.eye(_A)[y2] - p2[y1]; s2[_A + y1 * _A: _A + (y1 + 1) * _A] = blk
    return s1, s2
def _episodes(th_sample):                                         # list of (prob under sampler, y1, y2, R)
    p1, p2 = _probs(th_sample); out = []
    for y1 in range(_A):
        for y2 in range(_A):
            for r in (0, 1):
                out.append((p1[y1] * p2[y1, y2] * (_T[y1, y2] if r else 1 - _T[y1, y2]), y1, y2, r))
    return out
def _fd_grad(th, h=1e-6):
    g = _np.zeros_like(th)
    for i in range(len(th)):
        e = _np.zeros_like(th); e[i] = h; g[i] = (_J(th + e) - _J(th - e)) / (2 * h)
    return g
g_true = _fd_grad(th_cur)
pc1, pc2 = _probs(th_cur); ps1, ps2 = _probs(th_smp)
def _rho(y1, y2): return pc1[y1] / ps1[y1], pc2[y1, y2] / ps2[y1, y2]
# (a) on-policy score-function identity
g_on = sum(p * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_cur))
check("8a. E_pi[R (s1+s2)] = finite-difference grad J", float(_np.abs(g_on - g_true).max()), 1e-8)
# (b) sampler episodes, full sequence ratio rho1*rho2: exact
g_seq = sum(p * _rho(y1, y2)[0] * _rho(y1, y2)[1] * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_smp))
check("8b. E_mu[rho1 rho2 R (s1+s2)] = grad J exactly (sampler ratio, full sequence weight)", float(_np.abs(g_seq - g_true).max()), 1e-9)
# (c) token-level ratios only (what the objective uses): biased by the dropped prefix ratio; vanishes when sampler = policy
def _tok_grad(th_sample, clip=None, baseline=0.0):
    g = 0
    for p, y1, y2, r in _episodes(th_sample):
        q1, q2 = _rho(y1, y2) if th_sample is th_smp else (1.0, 1.0)
        if clip is not None: q1, q2 = min(q1, clip), min(q2, clip)
        s1, s2 = _score(th_cur, y1, y2); g = g + p * (r - baseline) * (q1 * s1 + q2 * s2)
    return g
g_tok = _tok_grad(th_smp)
bias_tok = float(_np.linalg.norm(g_tok - g_true)); gn = float(_np.linalg.norm(g_true))
print(f"   token-level sampler ratio, unclipped: ||bias|| = {bias_tok:.3e} (||grad J|| = {gn:.3e})")
check("8c. token-level weights are biased for a mismatched sampler (> 1e-4)", 0.0 if bias_tok > 1e-4 else 1.0, 0.0)
th_save = th_smp; th_smp = th_cur.copy()                          # sampler == policy: weights are 1
check("8c. ... and exact when the sampler equals the policy", float(_np.abs(sum(p * r * sum(_score(th_cur, y1, y2)) for p, y1, y2, r in _episodes(th_cur)) - g_true).max()), 1e-8)
th_smp = th_save
# (d) a constant baseline does not bias the token-level estimator, even under mismatch (E_mu[rho_t s_t | context] = 0)
b0 = _J(th_cur)
g_b = _tok_grad(th_smp, baseline=b0)
db = _np.abs(g_tok - g_b - 0).max()
check("8d. constant baseline b adds exactly 0 to the token-level sampler-ratio gradient: max |E_mu[sum_t rho_t b s_t]| over 20 coordinates", float(db), 1e-14)
# (e) clip bias bound with the sampler as the base measure: sum_a (pi - c mu)_+ <= TV(pi, mu) at every state
cc = 1.28
sts = [(pc1, ps1)] + [(pc2[y], ps2[y]) for y in range(_A)]
viol = max(float(_np.maximum(a - cc * b, 0).sum() - 0.5 * _np.abs(a - b).sum()) for a, b in sts)
check("8e. sum_a (pi - c mu)_+ <= TV(pi, mu) at all 5 states (c=1.28)", max(viol, 0.0), 0.0)
g_clip = _tok_grad(th_smp, clip=cc)
bound = 0.0
for p, y1, y2, r in _episodes(th_smp):
    q1, q2 = _rho(y1, y2); s1, s2 = _score(th_cur, y1, y2)
    bound += p * abs(r) * (max(q1 - cc, 0) * _np.linalg.norm(s1) + max(q2 - cc, 0) * _np.linalg.norm(s2))
bias_clip = float(_np.linalg.norm(g_clip - g_tok))
print(f"   clip c={cc}: ||g_clip - g_token-IS|| = {bias_clip:.3e} <= bound {bound:.3e}")
check("8e. clipping bias <= E_mu[sum_t (rho_t - c)_+ |A| ||s_t||]", max(bias_clip - bound, 0.0), 0.0)
# (f) outcome broadcast vs per-turn credit (on-policy, oracle baselines): both unbiased, different variance
eps_ = _episodes(th_cur)
V = sum(p * r for p, y1, y2, r in eps_)
Q = _np.array([sum(pc2[y1, y2] * _T[y1, y2] for y2 in range(_A)) for y1 in range(_A)])
def _moment(coef1, coef2):
    m1 = 0; m2 = 0
    for p, y1, y2, r in eps_:
        s1, s2 = _score(th_cur, y1, y2); g = coef1(y1, y2, r) * s1 + coef2(y1, y2, r) * s2
        m1 = m1 + p * g; m2 += p * float(g @ g)
    return m1, m2
gb_mean, gb_m2 = _moment(lambda y1, y2, r: r - V, lambda y1, y2, r: r - V)
gt_mean, gt_m2 = _moment(lambda y1, y2, r: Q[y1] - V, lambda y1, y2, r: r - Q[y1])
check("8f. broadcast estimator (R-V)(s1+s2) is unbiased for grad J", float(_np.abs(gb_mean - g_true).max()), 1e-8)
check("8f. per-turn estimator (Q-V) s1 + (R-Q) s2 is unbiased for grad J", float(_np.abs(gt_mean - g_true).max()), 1e-8)
v_b = gb_m2 - float(gb_mean @ gb_mean); v_t = gt_m2 - float(gt_mean @ gt_mean)
print(f"   total variance (trace): broadcast {v_b:.4f}, per-turn {v_t:.4f}, ratio {v_b/v_t:.3f}")
# turn-1 decomposition: E[(R-V)^2 ||s1||^2] = E[(Q-V)^2 ||s1||^2] + E[Var(R|y1) ||s1||^2]
lhs = sum(p * (r - V) ** 2 * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)
t1 = sum(p * (Q[y1] - V) ** 2 * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)
varR = _np.array([sum(pc2[y1, y2] * (_T[y1, y2] + 0) for y2 in range(_A)) - Q[y1] ** 2 for y1 in range(_A)])   # Var(R|y1) = E[R^2|y1] - Q^2 (R binary: E R^2 = E R)
t2 = sum(p * varR[y1] * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[0]) for p, y1, y2, r in eps_)
check("8f. turn-1 second moment: E[(R-V)^2||s1||^2] = E[(Q-V)^2||s1||^2] + E[Var(R|y1)||s1||^2]", abs(lhs - t1 - t2), 1e-13)
print(f"   turn-1: broadcast {lhs:.4f} = {t1:.4f} (credit) + {t2:.4f} (noise from turn 2)")
check("8f. martingale orthogonality: E[s1.s2] = 0 so E||s1+s2||^2 = E||s1||^2 + E||s2||^2", abs(sum(p * float(_score(th_cur, y1, y2)[0] @ _score(th_cur, y1, y2)[1]) for p, y1, y2, r in eps_)), 1e-15)
# (g) score second moment grows linearly with the number of policy tokens: L-step context-dependent chain, enumerated
Lmax = 5; Th = _rng.normal(0, 0.8, (Lmax, _A, _A)); P0 = _sm(_rng.normal(0, 0.8, _A))
def chain_moments(L):
    tot = 0.0; per = _np.zeros(L); cross = 0.0
    for seq in _np.ndindex(*([_A] * L)):
        pr = P0[seq[0]]; sc = []
        for t_ in range(L):
            prev = 0 if t_ == 0 else seq[t_ - 1]
            pt = P0 if t_ == 0 else _sm(Th[t_, prev])
            if t_ > 0: pr *= pt[seq[t_]]
            v = _np.zeros(L * _A * _A + _A)
            v[:_A] = (_np.eye(_A)[seq[0]] - P0) if t_ == 0 else 0
            if t_ > 0: v[_A + t_ * _A * _A + prev * _A: _A + t_ * _A * _A + (prev + 1) * _A] = _np.eye(_A)[seq[t_]] - pt
            sc.append(v)
        S_ = sum(sc); tot += pr * float(S_ @ S_)
        for t_ in range(L): per[t_] += pr * float(sc[t_] @ sc[t_])
    return tot, per
rows = []
for L in range(1, Lmax + 1):
    tot, per = chain_moments(L); rows.append((L, tot, per.sum()))
    print(f"   L={L}: E||sum_t s_t||^2 = {tot:.4f}   sum_t E||s_t||^2 = {per.sum():.4f}")
check("8g. E||sum_t s_t||^2 = sum_t E||s_t||^2 for L = 1..5 (score terms are martingale differences)", max(abs(a - b) for _, a, b in rows), 1e-12)
print("done; checks:", len(CHECKS))