# %% [markdown]
# # 0102 - Policy gradient from first principles, REINFORCE, baselines
# Companion to `notes/0102 - Policy Gradient and Baselines.pdf`.  Exact enumeration of the 121-sequence toy LM; fp64.
# %%
import sys, itertools
sys.path.insert(0, "..")
from _m1_lib import *
lm = TinyLM(V=4, L=4, seed=0, scale=1.0)
theta = lm.theta0.clone()
R = default_reward(lm)                       # binary verifier reward
p = lm.probs(theta)
S = score_matrix(lm, theta)                  # [N, P]  grad log pi(o)
P_ = theta.numel()
print("N =", lm.N, " P =", P_, " E[R] =", float(p @ R))
# %% [markdown]
# ## 1. Log-derivative trick:  grad J = E_pi[R * score]   (note Eq. 2)
# J(theta) = sum_o pi_theta(o) R(o).  Three independent computations of its gradient.
# %%
th = theta.clone().requires_grad_(True)
J = (lm.probs(th) * R).sum()
g_auto = torch.autograd.grad(J, th)[0]                       # differentiate the *definition* (through pi)
g_ld = (p[:, None] * R[:, None] * S).sum(0)                   # E_pi[R * score]  (the derived formula)
def Jfun(t): return float((lm.probs(t) * R).sum())
h = 1e-6
g_fd = torch.zeros_like(theta)
for i in range(P_):                                           # naive central finite differences
    e = torch.zeros_like(theta); e[i] = h
    g_fd[i] = (Jfun(theta + e) - Jfun(theta - e)) / (2 * h)
check("log-derivative formula == autograd of J", maxabs(g_ld, g_auto), 1e-14)
check("log-derivative formula == central finite differences", maxabs(g_ld, g_fd), 1e-8)
# %% [markdown]
# ## 2. Baselines are free:  E[b * score] = 0  for any b that does not depend on the sampled response
# %%
check("E_pi[score] = 0", float((p[:, None] * S).sum(0).abs().max()), 1e-14)
for b in [0.0, float(p @ R), 5.0, -3.0]:
    g_b = (p[:, None] * (R - b)[:, None] * S).sum(0)
    check(f"unbiased with constant baseline b={b:+.3f}", maxabs(g_b, g_auto), 1e-13)
# state-dependent baseline b_t(o_<t): arbitrary function of the prefix  (note Eq. for token-level baseline)
T = token_score_tensor(lm, theta)                            # [N, L, P]  token scores
gen = torch.Generator().manual_seed(5)
tab = {}
def bpref(prefix):
    if prefix not in tab: tab[prefix] = float(torch.randn(1, generator=gen)) * 3.0
    return tab[prefix]
B = torch.zeros(lm.N, lm.L)
for i, s in enumerate(lm.seqs):
    for t in range(len(s)): B[i, t] = bpref(s[:t])
g_state = (p[:, None] * ((R[:, None] - B)[:, :, None] * T).sum(1)).sum(0)
check("unbiased with arbitrary state-dependent baseline b(o_<t)", maxabs(g_state, g_auto), 1e-13)
# a baseline that depends on the token being scored breaks the proof (counter-example)
B_bad = 0.7 * (lm.tokens == 1).double() * lm.mask
g_bad = (p[:, None] * ((R[:, None] - B_bad)[:, :, None] * T).sum(1)).sum(0)
print("baseline depending on o_t: bias norm =", float((g_bad - g_auto).norm()))
assert float((g_bad - g_auto).norm()) > 1e-3
# and the extreme case b = R(o) (the sample's own reward) kills the gradient completely
g_self = (p[:, None] * (R - R)[:, None] * S).sum(0)
print("self-baseline gradient norm:", float(g_self.norm()), " true grad norm:", float(g_auto.norm()))
# %% [markdown]
# ## 3. Variance: exact quadratic in b, optimal baseline, and the "(0, 2 b*)" interval
# Var(b) = E||(R-b) s||^2 - ||grad J||^2  with  E||(R-b)s||^2 = E[R^2||s||^2] - 2b E[R||s||^2] + b^2 E||s||^2.
# %%
s2 = (S ** 2).sum(1)                                   # ||s||^2 per sequence
A0, A1, A2 = float((p * R ** 2 * s2).sum()), float((p * R * s2).sum()), float((p * s2).sum())
bstar = A1 / A2
def var_direct(b):
    g = (R - b)[:, None] * S                            # per-sequence estimator value
    return float((p * (g ** 2).sum(1)).sum() - (p[:, None] * g).sum(0).norm() ** 2)
def var_formula(b): return A0 - 2 * b * A1 + b * b * A2 - float(g_auto.norm() ** 2)
for b in [0.0, 0.3, 1.0, 2.0]:
    check(f"variance formula at b={b}", abs(var_direct(b) - var_formula(b)), 1e-13)
grid = torch.linspace(-2, 4, 6001)
vs = torch.tensor([var_formula(float(b)) for b in grid])
check("b* = E[R||s||^2]/E||s||^2 minimises variance (grid min vs V(b*))", max(0.0, var_formula(bstar) - float(vs.min())), 1e-12)
check("minimiser on grid is within grid step of b*", abs(float(grid[vs.argmin()]) - bstar), 2e-3)
EvR = float(p @ R)
print(f"b*={bstar:.4f}  E[R]={EvR:.4f}   Var(0)={var_formula(0.0):.4f}  Var(E[R])={var_formula(EvR):.4f}  Var(b*)={var_formula(bstar):.4f}")
# which side of E[R] does b* fall on?  b* - E[R] = Cov_pi(R, ||s||^2) / E||s||^2   (exact identity); sign = sign of the covariance
cov_Rs2 = float((p * R * s2).sum() - EvR * (p * s2).sum())
check("b* - E[R] == Cov_pi(R, ||s||^2) / E||s||^2  (exact identity)", abs((bstar - EvR) - cov_Rs2 / A2), 1e-14)
m1_ = float((p * R * s2).sum() / (p * R).sum()); m0_ = float((p * (1 - R) * s2).sum() / (p * (1 - R)).sum())
print(f"E[||s||^2 | R=1] = {m1_:.3f}   E[||s||^2 | R=0] = {m0_:.3f}   -> successes carry the larger score norm, so b* > E[R]")
assert (m1_ > m0_) == (bstar > EvR)
# interval claim: any b in (0, 2b*) has Var(b) < Var(0); b outside is worse
inside = all(var_formula(float(b)) < var_formula(0.0) for b in torch.linspace(0.01, 2 * bstar - 0.01, 50))
outside = all(var_formula(float(b)) > var_formula(0.0) for b in [2 * bstar + 0.01, 2 * bstar + 1.0, -0.01, -1.0])
assert inside and outside
check("Var(b) < Var(0) iff b in (0, 2 b*)  (50 inside, 4 outside points)", 0.0, 0.0)
# %% [markdown]
# ## 4. Naive vs efficient estimator (Monte-Carlo), and agreement with the exact gradient
# naive: one backward per sample, python loop.  efficient: ONE backward through the surrogate  -mean(sg(A) * log pi).
# %%
n = 2000
gen = torch.Generator().manual_seed(11)
idx = torch.multinomial(p, n, replacement=True, generator=gen)
b = EvR
# naive
th = theta.clone().requires_grad_(True)
acc = torch.zeros_like(theta)
for i in idx.tolist():
    lp = lm.seq_logprob(th)[i]
    acc += (float(R[i]) - b) * torch.autograd.grad(lp, th)[0]
g_naive = acc / n
# efficient
th = theta.clone().requires_grad_(True)
surrogate = ((R[idx] - b).detach() * lm.seq_logprob(th)[idx]).mean()    # gradient of this = REINFORCE estimate
g_eff = torch.autograd.grad(surrogate, th)[0]
check("naive per-sample loop == single backward through surrogate", maxabs(g_naive, g_eff), 1e-13)
# MC accuracy against the exact gradient
n2 = 400_000
idx2 = torch.multinomial(p, n2, replacement=True, generator=gen)
est = ((R[idx2] - b)[:, None] * S[idx2])
g_mc, se = est.mean(0), est.std(0) / n2 ** 0.5
ok_ = se > 0
z = float(((g_mc - g_auto).abs()[ok_] / se[ok_]).max())
check("MC REINFORCE gradient within 4 SE of exact (max z)", z, 4.0, kind="mc")
# %% [markdown]
# ## 5. Group baselines: leave-one-out is unbiased, "mean including self" is biased by (G-1)/G
# Exact enumeration of ALL G-tuples of samples (G=2: 14641 tuples, G=3: 1.77M tuples, evaluated per first sample).
# %%
N = lm.N
RR = R[None, :] ; pp = p[None, :]
def loo_and_mean(G):
    """exact E[(R_1 - baseline_1) s_1] for baseline = mean of others (LOO) and mean of all (incl. self)."""
    # sum over samples 2..G of rewards: build the joint over (o_2..o_G) as a tensor of (reward sum, prob)
    if G == 2:
        rs = R[None, :].expand(N, N); w = p[None, :].expand(N, N)
        sum_others_R = rs; prob_others = w
    elif G == 3:
        sum_others_R = (R[:, None] + R[None, :]).reshape(-1)[None, :].expand(N, -1)
        prob_others = (p[:, None] * p[None, :]).reshape(-1)[None, :].expand(N, -1)
        # N x N^2 : first sample index x joint of other two
    Ro1 = R[:, None]
    base_loo = sum_others_R / (G - 1)
    base_all = (Ro1 + sum_others_R) / G
    wloo = (prob_others * (Ro1 - base_loo)).sum(1)       # E over others of (R1 - b)
    wall = (prob_others * (Ro1 - base_all)).sum(1)
    g_loo = (p * wloo)[:, None] * S
    g_all = (p * wall)[:, None] * S
    return g_loo.sum(0), g_all.sum(0)
for G in (2, 3):
    gl, ga = loo_and_mean(G)
    check(f"G={G}: leave-one-out baseline is exactly unbiased", maxabs(gl, g_auto), 1e-13)
    check(f"G={G}: mean-including-self gives exactly (G-1)/G * grad J", maxabs(ga, (G - 1) / G * g_auto), 1e-13)
# %% [markdown]
# ## 6. Causality / reward-to-go (per-token rewards): both unbiased, reward-to-go has lower variance here
# Per-token reward r_t = 1[o_t == 1].  Full return G = sum_t r_t;  reward-to-go G_t = sum_{t'>=t} r_t'.
# %%
rt = (lm.tokens == 1).double() * lm.mask                    # [N, L]
Gfull = rt.sum(1)
Gto = rt.flip(1).cumsum(1).flip(1) * lm.mask                # reward-to-go
th = theta.clone().requires_grad_(True)
gJ = torch.autograd.grad((lm.probs(th) * Gfull).sum(), th)[0]
est_full = (T.sum(1) * Gfull[:, None])                       # sum_t score_t * G
est_rtg = (T * Gto[:, :, None]).sum(1)                       # sum_t score_t * G_t
g_full = (p[:, None] * est_full).sum(0); g_rtg = (p[:, None] * est_rtg).sum(0)
check("full-return estimator unbiased", maxabs(g_full, gJ), 1e-13)
check("reward-to-go estimator unbiased (dropped terms E[score_t r_t'] , t'<t, vanish)", maxabs(g_rtg, gJ), 1e-13)
v_full = float((p * (est_full ** 2).sum(1)).sum() - gJ.norm() ** 2)
v_rtg = float((p * (est_rtg ** 2).sum(1)).sum() - gJ.norm() ** 2)
print(f"trace-variance full return {v_full:.4f}   reward-to-go {v_rtg:.4f}")
assert v_rtg < v_full
# for an OUTCOME reward (one number at the end) reward-to-go == full return: nothing to drop
Gto_out = R[:, None] * lm.mask
check("outcome reward: reward-to-go == outcome at every t (no causality gain)", maxabs(Gto_out.sum(1), R * lm.lengths), 1e-15)
print("done; checks:", len(CHECKS))