| \documentclass[11pt]{article} |
| \usepackage{common} |
| \begin{document} |
| \notehead{0102}{Policy gradient from first principles, REINFORCE, baselines}{Policy-optimisation lineage, batch M1. Companion notebook: \texttt{notebooks/0102-policy-gradient-baselines.ipynb}.} |
|
|
| \begin{problem} |
| \question{How can we estimate $\nabla_\theta J$ from samples and scalar verifier rewards alone, and keep the variance low enough to train?} |
| We want to climb $J(\theta)=\E_{o\sim\pi_\theta}[R(o)]$. Two things stand in the way: the sampling distribution itself depends on $\theta$ (so we cannot push the gradient inside the expectation), and $R$ is a verifier verdict on a discrete string, which has no derivative with respect to anything. We need an estimator of $\nabla_\theta J$ that uses only samples and scalar rewards, and then a way to make that estimator usable (its variance, un-controlled, is what makes naive RL on long responses fail). |
| \end{problem} |
|
|
| \section{Notation} |
| \begin{center}\small\begin{tabular}{@{}ll@{}}\toprule |
| symbol & meaning\\\midrule |
| $J(\theta)=\sum_o\pi_\theta(o)R(o)$ & expected reward (prompt $q$ suppressed; a finite sum over responses)\\ |
| $s(o)=\nabla_\theta\log\pi_\theta(o)=\sum_t s_t$ & sequence score; $s_t=\nabla_\theta\log\pi_\theta(o_t\mid s_t)$ is the token score\\ |
| $G_t=\sum_{t'\ge t}r_{t'}$ & reward-to-go (for an outcome reward $G_t=R$ for all $t$)\\ |
| $b,\ b_t(o_{<t})$ & baseline: constant, resp.\ function of the prefix only\\ |
| $\hat g_b(o)=(R(o)-b)\,s(o)$ & single-sample gradient estimator\\ |
| $V(b)=\E\|\hat g_b\|^2-\|\nabla J\|^2$ & total variance (trace of the covariance)\\ |
| $G$ & group size (number of samples per prompt), used in Sec.~2.5\\\bottomrule\end{tabular}\end{center} |
|
|
| \section{Derivation} |
| \subsection{The log-derivative trick} |
| Assume $R$ does not depend on $\theta$ and $\pi_\theta(o)>0$ (true for a softmax, whose support is all of $\mathcal O_L$), so $\nabla\pi_\theta=\pi_\theta\nabla\log\pi_\theta$ holds everywhere. Then |
| \begin{align} |
| \nabla_\theta J&=\sum_o R(o)\,\nabla_\theta\pi_\theta(o)=\sum_o\pi_\theta(o)\,R(o)\,\nabla_\theta\log\pi_\theta(o)=\E_{o\sim\pi_\theta}\big[R(o)\,s(o)\big]\label{eq:ld}\\ |
| &=\E_{o\sim\pi_\theta}\Big[R(o)\sum_{t=1}^{|o|}\nabla_\theta\log\pi_\theta(o_t\mid s_t)\Big]\quad\text{(using 0101, Eq.~(2)).}\nonumber |
| \end{align} |
| The right side is an expectation under the \emph{current} policy of a quantity we can compute from samples: this is REINFORCE (Williams 1992, as used in arXiv:2510.26788 Eq.~2). No derivative of $R$ appears. |
| \ver{three computations of $\nabla J$ agree: autograd of $\sum_o\pi R$ vs.\ Eq.~\eqref{eq:ld}: $1.4\times10^{-17}$; vs.\ central finite differences: $3.4\times10^{-11}$.} |
|
|
| \subsection{The score has zero mean} |
| \begin{lemma}\label{lem:zero}For every state $s_t$: $\sum_a\pi_\theta(a\mid s_t)\nabla_\theta\log\pi_\theta(a\mid s_t)=0$, and hence $\E_{o\sim\pi_\theta}[s(o)]=0$.\end{lemma} |
| \emph{Proof.} $\sum_a\pi\nabla\log\pi=\sum_a\nabla\pi(a)=\nabla\sum_a\pi(a)=\nabla 1=0$. The sequence statement is the same computation on $\sum_o\pi(o)=1$. $\blacksquare$ |
| \ver{$\E_\pi[s]$ has max component $5.6\times10^{-17}$.} |
|
|
| \subsection{Baselines: which ones are free} |
| \paragraph{Constant (or prompt-only) baseline.} For $b$ independent of the sampled $o$: $\E[(R-b)s]=\E[Rs]-b\,\E[s]=\nabla J$ by Lemma~\ref{lem:zero}. Any $b$, even a wrong one, leaves the expectation unchanged. |
|
|
| \paragraph{State-dependent baseline.} Let $b_t=b_t(s_t)$ depend on the prefix only. Condition on $s_t$ (tower rule; the event $\{|o|\ge t\}$ that token $t$ exists is decided by $o_{<t}$, i.e.\ by $s_t$, so the random response length does not break the argument): |
| $\E[b_t(s_t)\,\nabla\log\pi(o_t\mid s_t)]=\E_{s_t}\big[b_t(s_t)\,\E_{o_t\mid s_t}\nabla\log\pi(o_t\mid s_t)\big]=0$ by Lemma~\ref{lem:zero}. This is the $B(q,o_{<t})$ of Dr.\ GRPO Appendix~A (arXiv:2503.20783, Eq.~5), with $B=V(s_t)$ the "value" that gives the advantage $A_t=G_t-V(s_t)$ (GAE, arXiv:1506.02438 Eq.~1, item 3 and Eqs.~2--3). |
| \ver{random $b(o_{<t})$ (values up to $\sim 3$): bias $1.1\times10^{-16}$; constants $b\in\{0,\,\E R,\,5,\,-3\}$: $\le2.1\times10^{-16}$.} |
|
|
| \paragraph{What is \emph{not} free.} The proof needs $b$ to be constant given the conditioning variable of the expectation $\E_{o_t\mid s_t}$. If $b$ depends on $o_t$ (or on anything after $t$) the factorisation fails. Extreme case: $b=R(o)$ makes $\hat g\equiv0$ although $\nabla J\neq0$. A baseline built from \emph{other} samples is fine because those are independent of $o_i$; a baseline that includes $o_i$ itself is not (Sec.~2.5). |
| \ver{$b=0.7\cdot\mathbb 1[o_t=1]$: bias norm $0.166$ ($\|\nabla J\|=0.182$); $b=R(o)$: estimator norm exactly $0$.} |
|
|
| \paragraph{Causality (reward-to-go).} Let $r_{t'}$ be a function of $o_{\le t'}$ (so of $s_t$ when $t'<t$). Then $\E[r_{t'}\,s_t]=\E_{s_t}[r_{t'}\E_{o_t|s_t}s_t]=0$ for $t'<t$, so these terms can be dropped: $\nabla J=\E\sum_t G_t\,s_t$ with $G_t=\sum_{t'\ge t}r_{t'}$. For a single outcome reward at the end $G_t=R$ for every $t$: there is nothing to drop and no variance is gained, which is why LLM-RL with verifier rewards relies on baselines, not on causality. |
| \ver{full return and reward-to-go are both unbiased to $4.2\times10^{-17}$; trace variance $2.86\to1.62$ on the per-token-reward toy.} |
|
|
| \subsection{Variance: why a baseline helps and exactly when} |
| By Sec.~2.3, $\E\hat g_b=\nabla J$ for every $b$, so the variance differs only through the second moment, which is a quadratic in $b$: |
| \begin{equation}V(b)=\underbrace{\E[R^2\|s\|^2]}_{A_0}-2b\underbrace{\E[R\|s\|^2]}_{A_1}+b^2\underbrace{\E\|s\|^2}_{A_2}-\|\nabla J\|^2.\label{eq:var}\end{equation} |
| Minimising: $b^\star=A_1/A_2=\E[R\|s\|^2]/\E[\|s\|^2]$, the mean reward \emph{weighted by the squared score norm}. Moreover $V(0)-V(b)=A_2\,b\,(2b^\star-b)$, so for $b^\star>0$: |
| \[V(b)<V(0)\iff 0<b<2b^\star.\] |
| Hence (i) a baseline is not automatically variance-reducing; (ii) the mean reward $b=\E R$ is optimal precisely when $\|s\|^2$ is uncorrelated with $R$ ($b^\star=\E R$), and otherwise only near-optimal. Intuition: $b^\star$ is the mean reward under weights $\|s\|^2$, and long, low-probability responses have the large $\|s\|$. If those responses are the \emph{successes} (rare successes are rare, hence low-probability), $b^\star$ sits \emph{above} the mean reward; if they are the failures it sits \emph{below} it. In general $b^\star>\E R\iff\Cov(R,\|s\|^2)>0$. |
| \ver{Eq.~\eqref{eq:var} vs.\ direct variance: $\le8.9\times10^{-16}$; on the toy $b^\star=0.302$, $\E R=0.166$, $V(0)=0.690$, $V(\E R)=0.516$, $V(b^\star)=0.472$; here the successes carry the larger score norm ($\E[\|s\|^2\mid R{=}1]=4.36$ vs.\ $\E[\|s\|^2\mid R{=}0]=2.00$), which is why $b^\star>\E R$; the interval statement holds at 50 interior and 4 exterior points.} |
|
|
| \subsection{Group baselines (preview of 0105)} |
| Draw $G$ i.i.d.\ responses $o_1,\dots,o_G$ for one prompt with rewards $R_i$. |
| \begin{itemize} |
| \item \emph{Leave-one-out (RLOO):} $b_i=\frac1{G-1}\sum_{j\ne i}R_j$ is independent of $o_i$, so $\E[(R_i-b_i)s_i]=\nabla J$. |
| \item \emph{Mean including self:} $\bar R=\frac1G\sum_jR_j$. Then $\E[(R_i-\bar R)s_i]=(1-\tfrac1G)\E[R_is_i]-\tfrac1G\sum_{j\ne i}\E[R_j]\E[s_i]=(1-\tfrac1G)\nabla J$ (independence and Lemma~\ref{lem:zero}). The estimator is unbiased \emph{up to the known scale} $(G-1)/G$, which a learning rate absorbs. Indeed $\frac{G}{G-1}(R_i-\bar R)=R_i-\frac1{G-1}\sum_{j\ne i}R_j$ (arXiv:2503.20783 App.~A). |
| \end{itemize} |
| \ver{exact enumeration of all $121^2$ ($G{=}2$) and $121^3$ ($G{=}3$) sample tuples: LOO error $\le2.1\times10^{-17}$; mean-including-self equals $\frac{G-1}{G}\nabla J$ to $\le1.0\times10^{-17}$.} |
|
|
| \subsection{Implementation: naive vs.\ efficient} |
| Naive: a Python loop calling \texttt{backward} on $(R_i-b)\log\pi(o_i)$ once per sample. Efficient: because $\hat g=\nabla_\theta[\sg(A)\,\log\pi_\theta]$ where $\sg$ is stop-gradient, one backward through the \emph{surrogate} $\frac1n\sum_i\sg(A_i)\log\pi_\theta(o_i)$ gives the whole estimate. (The surrogate is not $J$; its value is meaningless, only its gradient is.) The stop-gradient on the advantage is essential; later notes extend this idea to the ratio. |
| \ver{loop vs.\ single backward: $8.6\times10^{-16}$; Monte Carlo estimate from $4\times10^5$ samples within $2.0$ standard errors (max over components) of the exact gradient.} |
|
|
| \section{Boxed result} |
| \begin{result} |
| \[\boxed{\begin{aligned}\nabla_\theta J&=\E_{o\sim\pi_\theta}\Big[\sum_{t=1}^{|o|}\big(G_t-b_t(s_t)\big)\,\nabla_\theta\log\pi_\theta(o_t\mid s_t)\Big]\quad\text{for any }b_t=b_t(s_t),\\ |
| V(b)&=A_0-2bA_1+b^2A_2-\|\nabla J\|^2,\qquad b^\star=\frac{\E[R\|s\|^2]}{\E[\|s\|^2]},\qquad V(b)<V(0)\iff b\in(0,2b^\star)\end{aligned}}\] |
| \end{result} |
|
|
| \section{What breaks} |
| \begin{breaks}\small\begin{itemize} |
| \item \textbf{Off-policy samples.} Eq.~\eqref{eq:ld} is an expectation under the \emph{current} $\pi_\theta$. After one optimiser step the samples are from an old policy and the estimator is biased. \emph{Violated:} $o\sim\pi_\theta$. Fix: importance sampling (0103). |
| \item \textbf{Baseline sees the action.} Any $b$ that depends on $o_t$ or the future, including a batch mean that contains $o_i$ itself without the $G/(G{-}1)$ bookkeeping, changes the expectation. \emph{Violated:} $b\perp o_t\mid s_t$. |
| \item \textbf{Sparse rewards.} With success probability $p\ll1$ most batches contain only $R=0$ and the estimate is exactly zero or pure noise (0110). |
| \item \textbf{Variance grows with length.} $\|s\|^2=\|\sum_ts_t\|^2$ sums $|o|$ terms, so $A_0,A_2$ scale with response length; combined with an outcome reward the per-token credit is diluted. |
| \item \textbf{Reward depends on $\theta$.} If $R$ itself is learned and updated, there is an extra term $\E[\nabla R]$ missing from Eq.~\eqref{eq:ld}. |
| \item \textbf{No exploration beyond the support.} Responses the policy never samples never contribute; the estimator cannot raise a mode it does not see (the entropy theme of 0106). |
| \end{itemize}\end{breaks} |
|
|
| \nextnote{REINFORCE needs a fresh batch for every gradient step, which is wasteful when generation dominates cost. Note 0103 derives importance sampling to reuse a batch for several steps, and quantifies why the correction explodes for long sequences.} |
|
|
| \section*{References} |
| {\small\begin{itemize}[leftmargin=1.2em] |
| \item Qi et al., arXiv:2510.26788, Sec.~2, Eqs.~1--2 (objective and REINFORCE estimator) and Sec.~2.1, Eq.~5 (advantage $A=R-B(x)$ with a prompt-only baseline). |
| \item Liu et al., arXiv:2503.20783 (Dr.\ GRPO), Appendix A, Eqs.~4--6 (derivation, state-dependent baseline invariance, RLOO equivalence with factor $G/(G-1)$). |
| \item Schulman et al., \emph{High-Dimensional Continuous Control Using Generalized Advantage Estimation}, arXiv:1506.02438, Sec.~2, Eqs.~1--3 (list of equivalent $\Psi_t$, advantage). |
| \item REINFORCE is due to Williams (1992); it is cited here only as quoted in arXiv:2510.26788 and arXiv:2503.20783, not from the original. |
| \end{itemize}} |
| \end{document} |
|
|
