# %% [markdown]
# # 0106 - DAPO: each of its four techniques derived as the fix for a named pathology
# Companion to `notes/0106 - DAPO.pdf`.  fp64; exact enumeration where possible, Monte Carlo (flagged) elsewhere.
# %%
import sys, math
sys.path.insert(0, "..")
from _m1_lib import *
import numpy as np
# %% [markdown]
# ## 1. Clip-Higher (pathology: entropy collapse / exploration bottleneck)
# 1a. Asymmetric clip gradient table (eps_lo = 0.2, eps_hi = 0.28 as in DAPO Sec. 4.1) vs autograd; 1b. absolute probability headroom.
# %%
el, eh = 0.2, 0.28
def f_closed(r, A, el, eh): return A * min(r, 1 + eh) if A > 0 else A * max(r, 1 - el)
def grad_auto(r, A):
    lp = torch.tensor(math.log(r), requires_grad=True); rr = lp.exp()
    f = torch.minimum(rr * A, torch.clamp(rr, 1 - el, 1 + eh) * A)
    return float(torch.autograd.grad(f, lp)[0])
worst = 0.0; n = 0
for A in (-1.5, -0.2, 0.2, 1.5):
    for r in np.linspace(0.05, 3.0, 400):
        if min(abs(r - (1 - el)), abs(r - (1 + eh))) < 1e-6: continue
        live = (r < 1 + eh) if A > 0 else (r > 1 - el)
        worst = max(worst, abs(grad_auto(r, A) - (A * r if live else 0.0))); n += 1
check(f"asymmetric clip: autograd == table (A r on live set) on {n} points", worst, 1e-14)
print("   log-ratio limits: log(1+eps_hi)=%.4f  log(1-eps_lo)=%.4f  (symmetric eps=.2: +%.4f / %.4f)" % (math.log(1 + eh), math.log(1 - el), math.log(1.2), math.log(0.8)))
print("   absolute probability reachable by ONE clipped update of a token with pi_old = p (A > 0):")
print("      p        cap eps=0.20   cap eps=0.28    headroom(0.20)  headroom(0.28)")
for pp in (0.001, 0.01, 0.1, 0.5, 0.9):
    c1, c2 = min(1.0, 1.2 * pp), min(1.0, 1.28 * pp)
    print(f"     {pp:5.3f}     {c1:8.4f}      {c2:8.4f}        {c1-pp:8.4f}       {c2-pp:8.4f}")
check("DAPO example: 0.01 -> 0.012 at eps=0.2", abs(min(1, 1.2 * 0.01) - 0.012), 1e-15)
check("DAPO example: 0.9 -> min(1, 1.08) = 1.0 (no effective cap)", abs(min(1, 1.2 * 0.9) - 1.0), 1e-15)
# %% [markdown]
# 1c. Entropy dynamics of a softmax policy:  dH = -Cov_{a~pi}( log pi_a , dz_a )  (exact first order).  PG: dz = eta*pi*A;  NPG: dz = eta*A.
# %%
def entropy(z):
    lp = torch.log_softmax(z, 0); return float(-(lp.exp() * lp).sum())
def cov_formula(z, dz):
    pi = torch.softmax(z, 0); lg = pi.log()
    return float((pi * lg * dz).sum() - (pi * lg).sum() * (pi * dz).sum())
gen = torch.Generator().manual_seed(3)
worst_rel = 0.0
for trial in range(6):
    K = 5
    z = torch.randn(K, generator=gen) * 1.5
    A = torch.randn(K, generator=gen)
    pi = torch.softmax(z, 0); A = A - (pi * A).sum()                  # advantages must have zero pi-mean
    for name, dz in (("PG", pi * A), ("NPG", A)):
        eta = 1e-6
        dH_fd = (entropy(z + eta * dz) - entropy(z - eta * dz)) / (2 * eta)
        dH_cov = -cov_formula(z, dz)
        worst_rel = max(worst_rel, abs(dH_fd - dH_cov) / max(1e-12, abs(dH_cov)))
check("entropy change = -Cov(log pi, dz): finite difference vs formula (PG & NPG, 6 random policies)", worst_rel, 1e-6)
pi = torch.tensor([0.90, 0.05, 0.03, 0.02]); z = pi.log()
def dH_for(A):
    A = torch.tensor(A); A = A - (pi * A).sum(); return -cov_formula(z, pi * A)
print("   dH/d eta for tabular PG:  good action is the DOMINANT one (pi=.90): %+.5f   good action is a RARE one (pi=.02): %+.5f" %
      (dH_for([1., 0, 0, 0]), dH_for([0., 0, 0, 1.])))
assert dH_for([1., 0, 0, 0]) < 0 < dH_for([0., 0, 0, 1.])
check("positive-advantage dominant token lowers entropy; positive-advantage rare token raises it (signs)", 0.0, 0.0)
# %% [markdown]
# 1d. Exact multi-step dynamics on a 2-action bandit (rare correct action a, pi=0.01): how fast can p(a) grow per rollout iteration?
# Objective per inner step = sum_a pi_old(a) * f(r_a, A_a)  (exact expectation), tabular logits, plain gradient ascent.
# %%
def run_bandit(p0, eps_hi, outer=60, inner=400, lr=0.4, el=0.2):
    z = torch.log(torch.tensor([p0, 1 - p0])).clone()
    traj = [p0]
    for it in range(outer):
        pi_old = torch.softmax(z, 0).detach()
        A = torch.tensor([1.0, 0.0]) - pi_old[0]               # A_a = 1 - p, A_b = -p  (binary reward on action a)
        zz = z.clone().requires_grad_(True)
        for _ in range(inner):
            pi = torch.softmax(zz, 0); r = pi / pi_old
            f = torch.where(A > 0, A * torch.clamp(r, max=1 + eps_hi), A * torch.clamp(r, min=1 - el))
            obj = (pi_old * f).sum()
            g = torch.autograd.grad(obj, zz)[0]
            with torch.no_grad(): zz += lr * g
        z = zz.detach()
        traj.append(float(torch.softmax(z, 0)[0]))
    return torch.tensor(traj)
t20, t28, tinf = run_bandit(0.01, 0.2), run_bandit(0.01, 0.28), run_bandit(0.01, 10.0)
def first_cross(t, v=0.5):
    idx = (t >= v).nonzero()
    return int(idx[0]) if len(idx) else None
print("   iterations for p(a) to reach 0.5 starting from 0.01:  eps_hi=0.20 ->", first_cross(t20), "  eps_hi=0.28 ->", first_cross(t28), "  no upper clip ->", first_cross(tinf))
g20 = (t20[1:6] / t20[:5]); g28 = (t28[1:6] / t28[:5])
print("   per-iteration growth factor p'/p over first 5 iterations: eps=.20:", [round(float(x), 3) for x in g20], " eps=.28:", [round(float(x), 3) for x in g28])
print("   growth in iteration 1 slightly EXCEEDS 1+eps (1.251 vs 1.2; 1.342 vs 1.28): the clip stops only the DIRECT push on the rare action; the unclipped negative-advantage action keeps moving mass to it")
assert first_cross(t28) <= first_cross(t20)
check("larger eps_hi never slows the rare action (iterations to 0.5: eps_hi=.28 <= eps_hi=.20)", 0.0, 0.0)
# %% [markdown]
# ## 2. Dynamic sampling (pathology: zero-advantage groups).  P(dead group) = p^G + (1-p)^G
# %%
torch.manual_seed(8)                                   # torch.distributions samplers use the global RNG
gen = torch.Generator().manual_seed(8)
for G in (4, 16):
    for pp in (0.02, 0.3):
        ks = torch.distributions.Binomial(G, torch.tensor(pp)).sample((400000,))
        emp = float(((ks == 0) | (ks == G)).double().mean()); ex = pp ** G + (1 - pp) ** G
        se = math.sqrt(ex * (1 - ex) / 400000)
        check(f"P(dead) MC vs p^G+(1-p)^G  (G={G}, p={pp})  |diff|/SE", abs(emp - ex) / se, 4.0, kind="mc")
# expected live fraction for p ~ Beta(a,b):  1 - B(a+G,b)/B(a,b) - B(a,b+G)/B(a,b)
def lbeta(a, b): return math.lgamma(a) + math.lgamma(b) - math.lgamma(a + b)
def live_fraction(a, b, G): return 1 - math.exp(lbeta(a + G, b) - lbeta(a, b)) - math.exp(lbeta(a, b + G) - lbeta(a, b))
xs, ws = np.polynomial.legendre.leggauss(400); xs = 0.5 * (xs + 1); ws = 0.5 * ws
def live_quad(a, b, G):
    dens = np.exp(-lbeta(a, b)) * xs ** (a - 1) * (1 - xs) ** (b - 1)
    return float((ws * dens * (1 - xs ** G - (1 - xs) ** G)).sum())
for (a, b) in ((2, 5), (3, 3), (2, 8)):
    check(f"expected live fraction closed form vs quadrature, Beta({a},{b}), G=16", abs(live_fraction(a, b, 16) - live_quad(a, b, 16)), 1e-10)
a, b, G = 2.0, 8.0, 16
rho = live_fraction(a, b, G)
print(f"   Beta({a:.0f},{b:.0f}) prompts (mean p={a/(a+b):.2f}), G={G}: live fraction rho = {rho:.3f}  ->  oversampling factor 1/rho = {1/rho:.2f}")
# MC: prompts drawn until m live ones are found; expectation m / rho
torch.manual_seed(2)
gen = torch.Generator().manual_seed(2)
m_need, reps = 64, 600
draws = []
beta = torch.distributions.Beta(torch.tensor(a), torch.tensor(b))
for _ in range(reps):
    live = 0; drawn = 0
    while live < m_need:
        pp = beta.sample()
        k = torch.distributions.Binomial(G, pp).sample()
        drawn += 1; live += int(0 < k < G)
    draws.append(drawn)
d = torch.tensor(draws, dtype=torch.float64)
check("MC: prompts drawn to collect 64 live groups vs m/rho (z-score)", abs(float(d.mean()) - m_need / rho) / (float(d.std()) / reps ** 0.5), 4.0, kind="mc")
# Unbiasedness of dynamic sampling (exact, no MC): prompts q_j with success prob p_j and weights w_j; per-prompt Bernoulli-bandit gradient g = p_hat(1-p_hat)
# (centred advantage x score, see 0110).  Filtered batch: a prompt is accepted with probability P(live|q_j) (so accepted prompts have law w_j P(live_j)/rho) and its
# group contributes E[g | live, q_j] computed DIRECTLY from the binomial restricted to 0<k<G (not using g=0 on dead groups).
Gd = 16
pj = np.linspace(0.02, 0.98, 25); wj = pj ** (a - 1) * (1 - pj) ** (b - 1); wj /= wj.sum()
def binom_k(pp): return np.array([math.comb(Gd, k) * pp ** k * (1 - pp) ** (Gd - k) for k in range(Gd + 1)])
kk = np.arange(Gd + 1); gk = (kk / Gd) * (1 - kk / Gd)
Eg_all = np.array([float((binom_k(x) * gk).sum()) for x in pj])
Plive = np.array([float(binom_k(x)[1:Gd].sum()) for x in pj])
Eg_live = np.array([float((binom_k(x)[1:Gd] * gk[1:Gd]).sum()) / pl for x, pl in zip(pj, Plive)])
rho_j = float((wj * Plive).sum())
unfiltered = float((wj * Eg_all).sum())
filtered = float((wj * Plive / rho_j * Eg_live).sum())
check("dynamic sampling: E[g | accepted group] = (1/rho) * E[g | unfiltered group]  (exact, discrete prompt law)", abs(filtered - unfiltered / rho_j), 1e-15)
# the dilution identity: dividing by ALL samples instead of live samples scales the gradient by exactly the live fraction
print("   with global token normalisation, dead groups stay in the denominator: signal scales by rho =", round(rho, 3), "(identity, no estimation)")
# %% [markdown]
# ## 3. Token-level loss: global denominator is a RATIO estimator (bias O(1/#samples)); constant denominator is exactly unbiased
# Exact over all G-tuples (G=2,3) on the TinyLM.
# %%
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
theta = lm.theta0.clone(); R = default_reward(lm); p = lm.probs(theta); S = score_matrix(lm, theta); N = lm.N
lens = lm.lengths
grad_J = (p[:, None] * (R - float(p @ R))[:, None] * S).sum(0)
Lbar = float((p * lens).sum())
def exact_tokenlevel(G, normaliser):
    """E[ sum_i A_i * score(o_i) / D ] with centred advantages A_i = R_i - mean R; D = sum_i |o_i| (global) or G*Lbar (constant)."""
    if G == 2:
        R1, R2 = R[:, None].expand(N, N), R[None, :].expand(N, N)
        l1, l2 = lens[:, None].expand(N, N), lens[None, :].expand(N, N)
        P = p[:, None] * p[None, :]
        mu = (R1 + R2) / 2
        D = (l1 + l2) if normaliser == 'global' else torch.full((N, N), G * Lbar)
        w1 = P * (R1 - mu) / D; w2 = P * (R2 - mu) / D
        return w1.sum(1) @ S + w2.sum(0) @ S
    g = torch.zeros(S.shape[1])
    R2, R3 = R[:, None].expand(N, N), R[None, :].expand(N, N)
    l2, l3 = lens[:, None].expand(N, N), lens[None, :].expand(N, N)
    P23 = p[:, None] * p[None, :]
    for i in range(N):
        mu = (R[i] + R2 + R3) / 3
        D = (lens[i] + l2 + l3) if normaliser == 'global' else torch.full((N, N), G * Lbar)
        g += p[i] * ((P23 * (R[i] - mu) / D).sum()) * S[i]
        g += p[i] * ((P23 * (R2 - mu) / D).sum(1) @ S)
        g += p[i] * ((P23 * (R3 - mu) / D).sum(0) @ S)
    return g
for G in (2, 3):
    gc = exact_tokenlevel(G, 'const'); gg = exact_tokenlevel(G, 'global')
    check(f"G={G}: constant denominator G*E|o| gives (1-1/G) grad J / (G E|o|) * G exactly", maxabs(gc * Lbar, (1 - 1 / G) * grad_J), 1e-13)
    rel = float((gg - gc).norm() / gc.norm()); cs = float(gg @ grad_J / (gg.norm() * grad_J.norm()))
    print(f"   G={G}: global-token-mean estimator: relative deviation from the unbiased one = {rel:.4f}, cosine to grad J = {cs:.5f}")
# %% [markdown]
# ## 4. Overlong handling: soft punishment is a REWARD; loss masking is a reward equal to the baseline
# %%
def soft_overlong(Ly, Lmax, Lcache):
    if Ly <= Lmax - Lcache: return 0.0
    if Ly <= Lmax: return ((Lmax - Lcache) - Ly) / Lcache
    return -1.0
Lmax, Lcache = 16384, 4096
check("soft overlong: 0 at Lmax-Lcache", abs(soft_overlong(Lmax - Lcache, Lmax, Lcache)), 0.0)
check("soft overlong: -1 at Lmax (continuous with the truncation penalty)", abs(soft_overlong(Lmax, Lmax, Lcache) + 1.0), 0.0)
ls = np.arange(1, 24000)
vals = np.array([soft_overlong(int(x), Lmax, Lcache) for x in ls])
slope = float(np.abs(np.diff(vals[:Lmax])).max())          # largest per-token change of the ramp (up to Lmax)
check("soft overlong: ramp slope is exactly 1/Lcache (|max slope * Lcache - 1|)", abs(slope * Lcache - 1.0), 1e-9)
print("   hard truncation reward has a jump of size 1 at Lmax; the soft ramp spreads it over", Lcache, "tokens")
# masking == baseline reward  (TinyLM: 'truncated' = hit the length limit L=4 without EOS)
trunc = (lm.mask.sum(1) == lm.L) & (lm.tokens[:, -1] != lm.EOS)
b = float(p @ R)
g_masked = (p[:, None] * ((~trunc).double() * (R - b))[:, None] * S).sum(0)               # loss masked on truncated samples
Rprime = torch.where(trunc, torch.full_like(R, b), R)                                      # truncated sample gets reward = baseline
g_rb = (p[:, None] * (Rprime - b)[:, None] * S).sum(0)
check("loss-masking == assigning the baseline value as reward to truncated samples (exact)", maxabs(g_masked, g_rb), 1e-14)
th = theta.clone().requires_grad_(True)
Rlen = torch.where(trunc, torch.full_like(R, -1.0), torch.zeros_like(R)) + 0.0
g_total_reward = (p[:, None] * (R + Rlen - b)[:, None] * S).sum(0)
th = theta.clone().requires_grad_(True)
g_sep = torch.autograd.grad((lm.probs(th) * (R + Rlen)).sum(), th)[0]
check("shaping as reward: policy gradient of E[R + R_len] (separable objective)", maxabs(g_total_reward, g_sep), 1e-14)
g_hard = (p[:, None] * (R - b)[:, None] * S).sum(0)                                        # truncated treated as plain failures (R=0)
print(f"   ||g_masked - g_hard|| = {float((g_masked-g_hard).norm()):.4f}  vs ||grad J|| = {float(grad_J.norm()):.4f}  (filtering changes the objective)")
# %% [markdown]
# ## 5. Naive vs vectorised DAPO objective: asymmetric clip + global token-level normaliser + dynamic-sampling filter (gradient agreement)
# %%
th_old = lm.theta0.clone(); lq = lm.token_logprobs(th_old); q_old = lm.probs(th_old)
th_new = perturb(th_old, 1.0, seed=3)
gen = torch.Generator().manual_seed(17)
G, groups = 8, []
while len(groups) < 3:                                               # dynamic sampling: resample until the group is live (0 < k < G)
    idx = torch.multinomial(q_old, G, replacement=True, generator=gen)
    k = int(R[idx].sum())
    if 0 < k < G: groups.append(idx)
def dapo_naive(th):
    lp_all = lm.token_logprobs(th)
    tot, ntok = 0.0, 0
    for idx in groups:
        Rg = R[idx]; A = (Rg - Rg.mean()) / (Rg.std(unbiased=False) + 1e-6)
        for j, i in enumerate(idx.tolist()):
            for t in range(int(lm.lengths[i])):
                r = torch.exp(lp_all[i, t] - lq[i, t])
                a = A[j]
                term = a * torch.clamp(r, max=1 + eh) if a > 0 else a * torch.clamp(r, min=1 - el)
                tot = tot + term; ntok += 1
    return -tot / ntok
def dapo_vec(th):
    lp_all = lm.token_logprobs(th)
    idx = torch.cat(groups); Rg = R[idx].reshape(len(groups), G)
    A = ((Rg - Rg.mean(1, keepdim=True)) / (Rg.std(1, unbiased=False, keepdim=True) + 1e-6)).reshape(-1)
    r = (lp_all[idx] - lq[idx]).exp(); mk = lm.mask[idx]
    f = torch.where(A[:, None] > 0, A[:, None] * torch.clamp(r, max=1 + eh), A[:, None] * torch.clamp(r, min=1 - el))
    return -(f * mk).sum() / mk.sum()
t1 = th_new.clone().requires_grad_(True); t2 = th_new.clone().requires_grad_(True)
check("naive triple loop == vectorised DAPO loss (value)", abs(float(dapo_naive(t1).detach()) - float(dapo_vec(t2).detach())), 1e-13)
g1 = torch.autograd.grad(dapo_naive(t1), t1)[0]; g2 = torch.autograd.grad(dapo_vec(t2), t2)[0]
check("naive triple loop == vectorised DAPO loss (gradient)", maxabs(g1, g2), 1e-13)
print("done; checks:", len(CHECKS))