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