\documentclass[11pt]{article}
\usepackage{common}
\begin{document}
\notehead{0204}{Self-correction as a multi-turn MDP: masks and per-turn credit}{Batch M2. Companion notebook: \texttt{notebooks/0204-self-correction-multiturn-mdp.ipynb} (every episode and every group of sampled episodes enumerated exactly; fp64, executed).}

\begin{problem}
\question{Which tokens of a multi-turn transcript receive gradient, how is credit assigned across turns, and which baselines keep it unbiased?}
A prover that sees the Lean compiler's error message and revises its proof can solve problems that independent resampling does not; Goedel-Prover-V2 runs up to two revision rounds after the first attempt and reports a gain of about two points at pass@32 over its non-revising mode \cite{gps}. To \emph{train} such a model with RL, three things must be pinned down: (i) which tokens of the long multi-turn transcript receive gradient, (ii) how credit is assigned across turns, and (iii) which baselines keep the gradient unbiased. Getting any of these wrong gives a run that logs a healthy loss while optimising something else.
\end{problem}
\stepcounter{section}% the problem box counts as section 1: keeps every \S and notebook cross-reference stable

\section{Notation}
\begin{center}
\begin{tabular}{@{}p{0.36\linewidth}p{0.60\linewidth}@{}}
\toprule
Symbol & Meaning \\ \midrule
$T$ & maximum number of turns (initial attempt plus revisions) \\
$y^t$ & the model's output at turn $t$ (a token sequence) \\
$f^t=\Psi(x,y^t)$ & environment feedback (verifier verdict plus compiler messages); $\theta$-independent \\
$c^t=(x,y^1,f^1,\dots,y^{t-1},f^{t-1})$ & context of turn $t$ \\
$r_t=V(x,y^t)$, $\ R(\tau)$ & per-turn success; episode return (e.g.\ $\sum_tr_t$ or $\max_tr_t$) \\
$s_t=\nabla_\theta\log\pi_\theta(y^t|c^t)$ & turn-$t$ score; sums over the tokens of $y^t$ only \\
$m_\ell\in\{0,1\}$ & loss mask: $1$ for policy-sampled tokens, $0$ for the prompt, feedback and template tokens \\
$G$ & episodes per prompt in a group; $\bar R=\frac1G\sum_jR_j$, $\bar R_{-i}=\frac1{G-1}\sum_{j\ne i}R_j$ \\
$Q(c^t)=\E[R_{\ge t}|c^t]$ & value of a context ($R_{\ge t}=\sum_{s\ge t}r_s$) \\ \bottomrule
\end{tabular}
\end{center}

\section{Derivation}
\subsection{The trajectory and what the policy controls}
An episode is $\tau=(x,y^1,f^1,y^2,f^2,\dots,y^T)$, stopping early on the first success if the harness is configured so. Its probability is
\[
P_\theta(\tau)=D(x)\ \prod_{t=1}^{T}\Bigl[\pi_\theta(y^t\mid c^t)\ \Psi(f^t\mid x,y^t)\Bigr],\qquad \pi_\theta(y^t|c^t)=\prod_{\ell}\pi_\theta(y^t_\ell\mid c^t,y^t_{<\ell}).
\]
The prompt distribution $D$ and the verifier $\Psi$ do not depend on $\theta$. Hence
\begin{equation}\label{eq:mask}
\nabla_\theta\log P_\theta(\tau)=\sum_{t=1}^{T}\sum_\ell\nabla_\theta\log\pi_\theta(y^t_\ell|\cdot),
\end{equation}
containing \emph{only} policy-sampled tokens. With the score-function identity of note 0201, $\nabla J=\E[R(\tau)\sum_ts_t]$. The mask $m_\ell$ is therefore not an implementation convenience; it is the statement that the environment factors drop out of (\ref{eq:mask}). Notebook: enumerating all 9 episodes of a two-turn tabular environment, $\E[R(s_1+s_2)]$ equals the autograd $\nabla J$ to $8\times10^{-17}$.

\subsection{Unmasked feedback tokens bias the gradient}
If the loss also scores the feedback tokens, the LM (which assigns them probability through the same output softmax) contributes $R\,\nabla_\theta\log\pi_\theta(f^t|\cdot)$. Since $f^t$ is \emph{not} drawn from $\pi_\theta$, the zero-mean identity (Eq.~(6) of note 0201) does not apply: $\E[\nabla\log\pi_\theta(f)]\ne0$, and the term is a reward-weighted \emph{supervised} signal to imitate compiler output. In the toy (feedback token tied to the turn-2 head) the spurious term has norm $0.68$ against $\|\nabla J\|=0.48$. In an LM the contamination reaches every shared parameter.

\subsection{Per-turn credit: causality}
Take $s<t$. The reward $r_s=V(x,y^s)$ is a function of $c^t$. Conditional on $c^t$, $\E[s_t|c^t]=0$ (note 0201, Eq.~(6)). So $\E[r_s\,s_t]=\E[r_s\,\E[s_t|c^t]]=0$ and past rewards can be dropped from the turn-$t$ coefficient:
\begin{equation}\label{eq:rtg}
\nabla J=\E\Bigl[\sum_{t=1}^{T}R_{\ge t}\ s_t\Bigr],\qquad R_{\ge t}=\sum_{s\ge t}r_s .
\end{equation}
With a terminate-on-success harness and $R=\max_tr_t$ this is automatic (a turn that exists has no earlier success, so $R_{\ge t}=R$); with a sum of per-turn rewards it removes terms that have zero mean but non-zero variance (a variance reduction in the toy, not a general guarantee). Notebook ($R=r_1+r_2$): $\E[r_1s_2]=0$ to $10^{-17}$, both estimators equal $\nabla J$, and the single-episode total variance falls from $1.657$ to $0.947$.

\subsection{The unbiasedness condition for a baseline}
\begin{proposition}
Let $b_t$ be any random variable that is conditionally independent of $y^t$ given $c^t$. Then $\E[(R_{\ge t}-b_t)\,s_t]=\E[R_{\ge t}\,s_t]$.
\end{proposition}
\emph{Proof.} $\E[b_ts_t]=\E\bigl[\E[b_t|c^t]\ \E[s_t|c^t]\bigr]=0$. $\square$

Admissible: a constant; $b(x)$; the value $Q(c^t)$; any statistic of \emph{other} episodes (independent of episode $i$ altogether); the mean of \emph{siblings} that share $c^t$ but exclude sample $i$. Inadmissible (they depend on $y^t$ itself): the sample's own outcome, a group mean that includes the sample, a sibling mean that includes the sample. The variance-optimal baseline among functions of $c^t$ is (set the derivative of $\E[(R_{\ge t}-b)^2\|s_t\|^2|c^t]$ to $0$)
\[
b_t^\star(c^t)=\frac{\E[R_{\ge t}\|s_t\|^2\,|\,c^t]}{\E[\|s_t\|^2\,|\,c^t]}\ \approx\ Q(c^t).
\]

\subsection{The GRPO mean baseline is a rescaled leave-one-out baseline}
Both baselines are linear in the group. The identity
\[
R_i-\bar R=R_i-\tfrac1GR_i-\tfrac{G-1}G\bar R_{-i}=\tfrac{G-1}{G}\bigl(R_i-\bar R_{-i}\bigr)
\]
holds for every group. Leave-one-out is independent of episode $i$ so it is admissible and unbiased, hence the mean-including-self estimator has expectation exactly $\frac{G-1}{G}\nabla J$: the same direction, shrunk by $\frac{G-1}G$ (a $12.5\%$ shrink at $G=8$). The notebook enumerates all $9^4$ groups: the identity holds with maximal residual $5.6\times10^{-17}$, $\E[g_{\mathrm{mean}}]=\frac34\nabla J$ to $1.2\times10^{-15}$, and $\Var(g_{\mathrm{mean}})\cdot(G/(G-1))^2=\Var(g_{\mathrm{LOO}})$ to $7\times10^{-15}$. A baseline that uses the sample's own turn-2 outcome (here $b=r_2$) is biased: $\max|\E g-\nabla J|=0.079$.

(Dividing by the group standard deviation, as in GRPO, is a separate per-prompt reweighting; see note 0206 and \cite{drgrpo}. Dividing the loss by the response length, also in GRPO and in Goedel-Prover-V2's ``token-averaged'' loss (App.~E.1), adds a further length-dependent weight that this note does not analyse; \cite{drgrpo} Sec.~3.1 discusses both biases.)

\subsection{Episode-level versus per-turn credit}
With one scalar advantage $A=R-b(x)$ on every token of the episode, turn-1 tokens receive the credit for the luck of turn~2. Telescope:
\[
R-V(x)=\underbrace{\bigl[Q(y^1)-V(x)\bigr]}_{\text{turn 1}}+\underbrace{\bigl[R-Q(y^1)\bigr]}_{\text{turn 2}},
\]
where $V(x)=\E[R|x]$ and $Q(y^1)=\E[R|x,y^1]$. The correct per-turn advantages are the two brackets: the turn-1 gradient needs only $Q(y^1)-V(x)$, and $\E[(R-Q(y^1))\,s_1]=0$, so the second bracket added to the turn-1 coefficient is pure noise. Per-turn baselines remove it. Tree sampling gives them unbiasedly: draw $G_1$ first attempts, $G_2$ revisions for each; for turn $2$ use the sibling leave-one-out mean (siblings share $c^2$), for turn~1 use the other branches' mean value $\hat Q_a-\frac1{G_1-1}\sum_{a'\ne a}\hat Q_{a'}$. The notebook verifies $\E[g_{\mathrm{tree}}]=\nabla J$ to $2\times10^{-16}$. Whether it has lower variance at a fixed episode budget is an empirical question, and in the toy the answer is no: with $4$ episodes, trace variance is $0.414$ (no baseline), $0.237$ (reward-to-go, no baseline), $0.252$ (flat LOO), $0.486$ (tree $2\times2$). The sibling baseline estimates $Q(c^2)$ from $G_2-1$ samples, inflating the variance of $(R-b)$ by $\tfrac{G_2}{G_2-1}$ relative to the ideal baseline; with $G_2=2$ that costs a factor $2$ and outweighs the conditioning gain. A constant offset of $+3$ on all returns multiplies the no-baseline variance twelve-fold ($0.41\to4.91$) and leaves both LOO variants unchanged. Per-turn baselines pay off only when $G_2$ is moderate or a learned $Q(c^t)$ is available (cf.\ VinePPO's Monte-Carlo value estimates \cite{vineppo}).

\subsection{Training only the correction round on stale contexts}
Goedel-Prover-V2's RL phase uses a 50\% / 50\% multi-task mix of whole-proof inputs and \emph{first-round self-correction} inputs, one epoch, GRPO-style \cite{gps}. Each self-correction input is a fixed triple (statement, an output previously generated by the SFT model, the compiler error message) collected before RL (App.~E.1), and the authors state that they tried tool-use and multi-turn RL, found them immature, and chose this single-turn formulation instead (App.~E.2). So the contexts really are drawn from a frozen distribution, as modelled below. Write the full gradient for two turns, $\nabla J=\E[Q(y^1)\,s_1]+\E_{y^1\sim\pi_1}\E_{y^2\sim\pi_2}[R\,s_2]$, and the second term equals $\E_{y^1\sim\pi_1}\nabla_{\theta}Q(y^1)$ when only the turn-2 head is varied. Training the revision round alone on contexts drawn from a frozen distribution $\mu$ estimates
\[
g_{\mathrm{trunc}}=\sum_a\mu(a)\,\nabla_{\theta_2}Q(a)
\]
(notebook: closed form and enumeration agree to $9\times10^{-18}$). It differs from the $\theta_2$-block of $\nabla J$ by $\sum_a(\mu-\pi_1)(a)\nabla Q(a)$ (norm $0.104$ in the toy with $\mu=(0.6,0.3,0.1)$) and omits the turn-1 block entirely (norm $0.458$). The whole-proof half of the mix supplies a bandit gradient for turn~1, so the two halves together approximate the full gradient when $\mu\approx\pi_1$ and the mixture weights match the desired turn weighting. The approximation degrades as $\pi_1$ drifts from the policy that produced the contexts.

\section{Boxed result}
\begin{result}
\[
\nabla J=\E\Bigl[\sum_{t=1}^{T}\bigl(R_{\ge t}-b_t\bigr)\sum_{\ell}m_\ell\,\nabla_\theta\log\pi_\theta(y^t_\ell|\cdot)\Bigr],\qquad m_\ell=0\ \text{on prompt, feedback and template tokens},
\]
unbiased for every $b_t$ with $b_t\perp y^t\mid c^t$; the group-mean baseline (sum-over-tokens loss, no $1/|y|$ or $1/\hat\sigma$ factors) gives $\frac{G-1}G\nabla J$ exactly; and
$R-V(x)=[Q(y^1)-V(x)]+[R-Q(y^1)]$ splits the episode advantage into the two per-turn advantages.
\end{result}

\section{What breaks}
\begin{breaks}
\begin{tabular}{@{}p{0.29\linewidth}p{0.29\linewidth}p{0.36\linewidth}@{}}
\toprule
Failure & Assumption violated & Consequence \\ \midrule
Loss on feedback or template tokens & $\theta$-independent environment tokens are masked & Adds $\E[R\nabla\log\pi(f)]\ne0$; model learns to imitate compiler output (norm $0.68$ vs $0.48$ in the toy) \\
Baseline uses own outcome / group mean with self & $b_t\perp y^t\mid c^t$ & Bias; $\frac{G-1}G$ shrink for the mean-with-self baseline, a direction error for outcome-dependent ones \\
Per-prompt std normalisation & uniform per-prompt weight & Reweights prompts by $1/\sigma_x$; see 0206 \\
Stale contexts $\mu\ne\pi_1$ & on-policy contexts & Bias $\sum(\mu-\pi_1)\nabla Q$; missing turn-1 term when only the correction round is trained \\
Uninformative feedback & $\Psi$ reveals the error & Revision $\approx$ resampling; Goedel's ablation: removing compiler messages hurts markedly, removing prior chain-of-thought hurts slightly \cite{gps} \\
Unshaped multi-turn RL & turn-2 gradient rewards changing a wrong answer & Policy learns to rarely change its answer and turn-2 responses couple tightly to turn-1; SCoRe adds a Stage~I initialisation and reward shaping toward progress \cite{score} \\
Small $G_2$ sibling baselines & baseline noise negligible & Variance inflation $G_2/(G_2-1)$; can lose to a flat leave-one-out at equal budget (toy: $0.49$ vs $0.25$) \\
Context growth over turns & fixed context budget & Later turns see long prefixes (cost, truncation); a training-inference mismatch if the context is truncated differently at test time \\ \bottomrule
\end{tabular}
\end{breaks}

\nextnote{Note 0205 turns to weight-space averaging: after SFT/RL pass@1 is higher and diversity lower. Goedel-Prover-V2 interpolates the trained weights back toward the base model and reports that pass@$N$ rises again; we analyse that interpolation to first order.}

\begin{thebibliography}{9}
\bibitem{gps} Y.~Lin et al. Goedel-Prover-V2. arXiv:2508.03613, Sec.~2.3 (RL: 50\% whole-proof / 50\% first-round self-correction inputs, hybrid GRPO), Sec.~3.2 (evaluation: two additional self-correction rounds, 40K total tokens), Sec.~3.3 (about $+2$ points at pass@32 from self-correction), Sec.~3.5 (compiler-message ablation), App.~E.1--E.2 (RL data and why multi-turn RL was not used).
\bibitem{score} A.~Kumar et al. Training Language Models to Self-Correct via Reinforcement Learning. arXiv:2409.12917, Sec.~5.1--5.2 (Stage~I initialisation, Stage~II with reward shaping), Fig.~6 (collapse of multi-turn RL).
\bibitem{rloo} A.~Ahmadian et al. Back to Basics: Revisiting REINFORCE Style Optimization. arXiv:2402.14740, Sec.~2.3 (leave-one-out baseline).
\bibitem{vineppo} A.~Kazemnejad et al. VinePPO. arXiv:2410.01679 (Monte-Carlo value estimates for credit assignment).
\bibitem{drgrpo} Z.~Liu et al. Understanding R1-Zero-Like Training: A Critical Perspective. arXiv:2503.20783 (std and length normalisation bias in GRPO).
\bibitem{dsmath} Z.~Shao et al. DeepSeekMath. arXiv:2402.03300, Sec.~4.1 (GRPO group baseline).
\end{thebibliography}
\end{document}