Revisiting RL for LLM Post-Training - Sample-Efficient Algorithms

This post reviews PPO and related policy-gradient methods for large language models from four perspectives.

For application-driven objectives, see Applications in Continual Learning.

Need a refresher on PPO? Start with the basics.

Learning without value networks

FlashREINFORCE and SPO use statistical baselines; CoPE adds action forecasting to GRPO. All three learn without a value network.

FlashREINFORCE

• No learned critic: For a batch of $B$ independent prompts, center the observed rewards:

  \[\bar R=\frac{1}{B}\sum_{j=1}^{B}R_j, \qquad A_i=R_i-\bar R.\]

  A response above the batch mean receives positive advantage; one below it receives negative advantage.

• Trajectory filtering: Weight token updates by current-to-rollout probability ratios. For each sampled token, let $p_t$ and $q_t$ be its rollout and current-policy probabilities at the same history:

  \[\begin{aligned} d_t &= p_t\log\frac{p_t}{q_t}+(1-p_t)\log\frac{1-p_t}{1-q_t},\\ \bar D &= \frac{1}{T}\sum_{t=1}^{T}d_t,\qquad m=\mathbf{1}[\bar D\leq\delta]. \end{aligned}\]

  Keep the whole trajectory only when $m=1$; $T$ counts model-generated tokens and $\delta$ is the threshold. This Bernoulli KL is a lower-bound proxy for full-vocabulary KL. Details.

• Equal trajectory weight: For token losses $\ell_{i,t}$ that include advantage and importance weighting, aggregate as

  \[L=\frac{1}{B}\sum_{i=1}^{B}m_i \left(\frac{1}{T_i}\sum_{t=1}^{T_i}\ell_{i,t}\right).\]

  Here $m_i$ is the trajectory gate above. Averaging within each response removes its explicit length multiplier. Each fresh batch receives one update.

SPO

• Persistent baseline: For prompt $x$, compare reward $R$ with the tracker’s estimate before this rollout:

  \[A=R-\hat v_{\mathrm{old}}(x), \qquad \tilde A=\frac{A-\bar A_{\mathcal B}}{\sigma_{A,\mathcal B}}.\]

  The baseline uses past outcomes; normalization uses the current batch $\mathcal B$ across prompts.

• Adaptive memory: The Bayesian tracker has an equivalent moving-average update:

  \[\hat v_{\mathrm{new}}(x)=\hat v_{\mathrm{old}}(x) +\eta_x\bigl(R-\hat v_{\mathrm{old}}(x)\bigr).\]

  The step size $\eta_x$ reflects policy change and accumulated evidence. A larger policy shift discounts older evidence more strongly. The tracker also supports prioritized prompt sampling.

CoPE

• Forecast successful actions: From history $h_t$, predict the recorded action chunk $a_{t:t+K}$. For the successful subset $\mathcal S_x$ of a rollout group:

  \[L_{\mathrm{AF},x}= -\frac{\displaystyle\sum_{\tau\in\mathcal S_x}\sum_t \log p_\theta(a_{t:t+K}\mid h_t)} {\max(1,|\mathcal S_x|)}.\]

  This averages supervision over successful trajectories; a group with no success contributes zero forecasting loss.

• Shared policy backbone: Minimize

  \[L=L_{\mathrm{GRPO}}+\eta\,\mathbb E_x[L_{\mathrm{AF},x}].\]

  The coefficient $\eta$ controls the extra supervision. Acting and forecasting use the same weights with different instructions; inference still takes one action at a time.

Changing what critics do

Use value estimates to learn before a long rollout finishes.

TEMPO — dots

Test-Time-Scaled Value Estimation with Macro-Step Policy Optimization, from Xiaohongshu’s dots team.

• Macro-step returns: Resume saved states for $H$ interactions. In undiscounted notation, branch $i$ supplies

  \[G_i=\sum_{h=0}^{H-1}r_{i,h}+\hat V(s_{i,H}), \qquad y(s)=\frac{1}{n}\sum_{i=1}^{n}G_i.\]

  Terminal states have zero remaining value. $G_i$ trains the actor; $y(s)$ targets the starting-state value.

• Generative critic: Reason over history, optionally using tools and privileged environment information. Call it at segment boundaries to limit cost.

• Learn to evaluate: Warm up with complete-trajectory returns, then use TD targets. Reward critic estimates by normalized accuracy and center rewards within each state’s group. Actor and critic share parameters and use GRPO, with different prompts and rewards.

The appendix derives importance correction for reused prefixes; segment updates alone do not remove policy staleness.

AC2

• Short continuations: From replayed prefix $s$, sample chunks $c_i$. Once the problem passes critic-readiness checks, unfinished chunks can use predicted endpoint values:

  \[v_i=\begin{cases} R_i, & \text{complete response},\\ V_\theta(s\cdot c_i), & \text{unfinished chunk}. \end{cases}\]

  The symbol $\cdot$ means concatenation. A trusted prediction can replace the cost of completing a response.

• Chunk advantages: Compare $G$ continuations of the same prefix:

  \[\bar v=\frac{1}{G}\sum_{j=1}^{G}v_j, \qquad \hat A_i=v_i-\bar v.\]

  All new tokens in chunk $c_i$ share $\hat A_i$. Train the shared generative critic at $s$ to output the rounded group mean $\bar v$ as a value token; replayed prefix tokens receive no actor loss.

• Full-rollout audits: Complete a fraction of ready problems to check critic predictions against actual outcomes.

Changing what critics see

Give the critic additional evidence for estimating value and assigning credit.

BPCO

• Ground the critic: With outcome-only rewards and no discounting, use

  \[R_{\min}\leq V_\phi(s_t)\leq R_{\max}, \qquad L_V=\mathbb E[(V_\phi(s_t)-R)^2].\]

  Bound the value head and fit observed final rewards, avoiding targets built from the critic’s own predictions. Optional reference answers or rubrics enter only the critic.

• Preserve the learning signal: For a response of length $L$, use

  \[\lambda_\pi(L)=1-\frac{1}{\alpha L}, \qquad \lambda_V=1.\]

  Here $\alpha>0$ controls GAE decay. Longer responses use slower decay, preserving distant outcome feedback. Keep raw GAE advantages instead of dividing by the batch standard deviation.

• Divergence PPO: If $p_t$ is the sampled token’s rollout probability, replace the fixed ratio-clipping interval with

  \[\left[1-\frac{\epsilon}{p_t},\;1+\frac{\epsilon}{p_t}\right].\]

  Its boundaries correspond to an absolute probability change of $\epsilon$. Clipping limits the incentive to move farther, rather than enforcing a hard constraint.

LVRL

• Frozen LLM critic: Read the completed trajectory, its outcome, and related rollouts to estimate intermediate values $\tilde V_t$. These are hindsight assessments, not prefix-only predictions; no critic training is needed.

• Anchored credit: For group mean reward $\bar R_G$ and trajectory reward $R$, fix the endpoints and take differences:

  \[\begin{aligned} \tilde V_0&=\bar R_G,\qquad \tilde V_T=R,\\ \bar A_t&=\tilde V_{t+1}-\tilde V_t,\\ \sum_{t=0}^{T-1}\bar A_t&=R-\bar R_G. \end{aligned}\]

  Actual outcomes fix the total credit; the LLM decides its distribution across actions. Divide $\bar A_t$ by the group-wide step standard deviation for training. All tokens in one action share its advantage; anchoring does not guarantee correct local attribution.

Stabilizing policy updates

Control errors from asynchronous sampling, mismatched probabilities, and noisy critic training.

SAO

• Single-rollout sampling: Collect one trajectory per prompt asynchronously. A learned critic supplies the baseline, so training need not wait for sibling rollouts.

• Token masking: Let $\mu$ be the logged rollout policy and $\pi_\theta$ the learner:

  \[\begin{aligned} \rho_t&=\frac{\pi_\theta(a_t\mid s_t)}{\mu(a_t\mid s_t)},\\ w_t&=\rho_t\,\mathbf 1[1-\epsilon_\ell<\rho_t<1+\epsilon_h]. \end{aligned}\]

  Use $w_t$ to weight the token’s advantage-gradient contribution. Outside the interval its contribution is zero, regardless of advantage sign.

• Skip-observation GAE: Let $t^+$ index the next model-generated token, skipping intervening observation tokens:

  \[\begin{aligned} \delta_t&=r_t+\gamma V(s_{t^+})-V(s_t),\\ \hat A_t&=\delta_t+\gamma\lambda\hat A_{t^+}. \end{aligned}\]

  The next state still includes tool feedback in its context. GAE simply does not recurse through each tool-output token; terminal continuation values are zero.

• Critic updates: Warm up the critic, freeze its attention layers, and update it twice per actor update to help it track the changing policy.

Score Centering

• Why center the score? At a fixed prefix $s$, let $p_a=\pi_\theta(a\mid s)$ be the learner, $q_a=\mu(a\mid s)$ the sampler, and $\psi_a=\nabla_\theta\log p_a$ the score. Write $A=\hat A(s,a)$ for the advantage, held fixed during the actor update. Kunhao Zheng’s decomposition makes the unwanted drift explicit:

  \[\mathbb E_q[A\psi] =\underbrace{\operatorname{Cov}_q(A,\psi)}_{\text{signal}} +\underbrace{\mathbb E_q[A]\,\mathbb E_q[\psi]}_{\text{drift}}.\]

  On-policy, $\mathbb E_p[\psi]=0$. When sampling and training distributions differ, that identity need not hold under $q$: even adding a constant to every advantage can change the expected update.

• Remove that drift: Subtract the sampler’s expected score:

  \[\begin{aligned} c(s)&=\sum_a q_a\psi_a,\\ g_{\mathrm{SC}}(s)&=\mathbb E_q\!\left[A\bigl(\psi-c(s)\bigr)\right] =\operatorname{Cov}_q(A,\psi). \end{aligned}\]

  Centering advantages by their exact mean at this prefix would zero the first factor of the drift term; score centering zeros the second. A batch-wide zero mean does not guarantee a zero mean at every prefix. The remaining covariance is still under $q$, so this is not a complete off-policy correction.

• Practical correction: Store sampler probabilities for its top-$k$ tokens $H$ and approximate the remaining mass using the trainer’s distribution:

  \[\begin{aligned} \rho&=\frac{1-\sum_{a\in H}q_a}{1-\sum_{a\in H}p_a},\\ \hat c(s)&=\sum_{a\in H}(q_a-\rho p_a)\psi_a. \end{aligned}\]

  This assigns unlogged tokens probability $\hat q_a=\rho p_a$ and uses $\sum_a p_a\psi_a=0$ to avoid summing over the tail. Hold the correction coefficients fixed when differentiating the loss. With importance weights, center $w_a\psi_a$ by subtracting $\mathbb E_{a\sim q}[w_a\psi_a]$ instead.

EasyPPO

• Actor-only truncation filtering: Let $C_i=1$ if response $i$ finishes within the length limit. Apply $C_i$ only to actor loss; keep every rollout in critic training. The critic then learns $\mathbb E[R\mid s]$, rather than $\mathbb E[R\mid s,C=1]$.

• Noise-weighted critic loss: Let $\hat\sigma_x$ be the return standard deviation across rollouts of prompt $x$:

  \[\begin{aligned} w_x&\propto\frac{1}{\max(\hat\sigma_x,\varepsilon)},\\ L_V&=\mathbb E_{x,i,t}\!\left[w_x\bigl(V_\phi(s_{i,t})-R_i\bigr)^2\right]. \end{aligned}\]

  High-variance prompts receive less weight, even when their mean reward is the same. This changes their influence on training, not their reward targets; $\varepsilon$ prevents division by zero.

• Smaller critic minibatches: Clip each minibatch gradient $g$ separately:

  \[g_{\mathrm{clip}}=g\min\!\left(1,\frac{c}{\lVert g\rVert}\right).\]

  The threshold $c$ limits gradient norm. Smaller minibatches restrict how many other samples an outlier can suppress through this shared rescaling. The actor retains the PPO objective.

StableReinforce

Martin Marek’s October 25 thread identifies StableReinforce as REINFORCE with Monte Carlo score centering. His custom loss reproduces its gradients within numerical precision in a controlled probe. This is independent reverse-engineering evidence; Tinker has not published the full policy-gradient formula.

• Estimate the mean score by sampling: Marek infers eight auxiliary next-token draws at each prefix:

  \[\begin{aligned} u_k&\sim q(\cdot\mid s),\qquad K=8,\\ \hat c(s)&=\frac1K\sum_{k=1}^{K}\nabla_\theta\log p_\theta(u_k\mid s),\\ \hat g_{\mathrm{SC}}&=A\left[\nabla_\theta\log p_\theta(a\mid s)-\hat c(s)\right]. \end{aligned}\]

  With independent sampler draws, $\hat c$ is an unbiased estimate of $c(s)$. This trades the original paper’s deterministic top-$k$ approximation for Monte Carlo noise; it does not make the whole update an unbiased on-policy gradient.

• Why the gradient gives it away: For logits $z$ and a one-hot token vector $e_a$, $\nabla_z\log p(a)=e_a-p$. Subtracting the sampled mean cancels the dense $-p$ term:

  \[\nabla_z\log p(a)-\frac1K\sum_k\nabla_z\log p(u_k) =e_a-\frac1K\sum_k e_{u_k}.\]

  Only the selected tokens remain. Marek observes this sparse pattern in StableReinforce, then matches gradients across all layers with a custom centered loss. The cancellation describes the policy-gradient term; a separate KL penalty can add a dense gradient.

• Official API and KL: Tinker accepts caller-supplied advantages and records sampling information for stability. Its default, optional $k_1$ regularizer is

  \[L_{\mathrm{KL}}=\beta\sum_t\left[\log q(a_t\mid s_t)-\log p_\theta(a_t\mid s_t)\right].\]

  This estimates $\beta\,\mathrm{KL}(q|p_\theta)$ at sampled prefixes, pulling the learner toward the sampler rather than a frozen reference model. The documentation also offers a $k_3$ estimator. Score centering changes the reward-weighted update; KL separately limits divergence.

Details

Agent: FlashREINFORCE, SPO, TEMPO, LVRL, SAO, EasyPPO, CoPE.

LLM: FlashREINFORCE, AC2, BPCO, Score Centering, EasyPPO.

AC2’s chunks divide a proof into segments, without external interaction.

See details on the tasks →