Skip to main content
Every method here is one loss over the tokens a model wrote. For the policy-gradient family the loss is L(θ)=−∑i∑tci,t log⁡πθ(ai,t∣hi,t),ci,t held constant,\mathcal{L}(\theta) = -\sum_{i}\sum_{t} c_{i,t}\,\log \pi_\theta(a_{i,t} \mid h_{i,t}), \qquad c_{i,t} \text{ held constant}, so the gradient step is ∑i,tci,t ∇θlog⁡πθ\sum_{i,t} c_{i,t}\,\nabla_\theta \log \pi_\theta, and a method is a choice of the per-token coefficient cc. The others (SFT, DPO, a reward model, distillation) are a different loss, written out the same way. The plain-words course is Learn; the numbers behind every default are in whileai/simulations/defaults.py with their sources; equation numbers are the papers’. Each method’s own page says where it runs: the hosted trainer (wai.train), prime-rl (wai.prime_rl_config), or your own loop (method.update).

​
Notation

A batch holds BB trajectories. Trajectory ii answers prompt xix_i with reward RiR_i and TiT_i action tokens ai,1…ai,Tia_{i,1} \ldots a_{i,T_i}, each with history hi,th_{i,t} (the prompt and the tokens before it). Two policies can see every token: μ\mu, the policy that wrote it (a row’s behavior_logprobs; the rollout engine, or last week’s served model), and πθ\pi_\theta, the one being trained (logprobs). Their per-token ratio is ρi,t=πθ(ai,t∣hi,t)μ(ai,t∣hi,t)=exp⁡(logprobs[t]−behavior_logprobs[t]),\rho_{i,t} = \frac{\pi_\theta(a_{i,t} \mid h_{i,t})}{\mu(a_{i,t} \mid h_{i,t})} = \exp\big(\texttt{logprobs}[t] - \texttt{behavior\_logprobs}[t]\big), and πref\pi_{\text{ref}} is a frozen reference (the starting weights) where a method keeps one. A token the environment wrote (tool output, an observation) has action_mask false: it is not one of the TiT_i, gets c=0c = 0, and the critic methods bootstrap over it. When behavior_logprobs is absent the row is taken as on-policy, ρ=1\rho = 1.

​
The hosted four: wai.train(method=)

SFT. Maximum likelihood on the rows as pushed, over the assistant tokens the loss mask keeps (the trainer does not filter on reward, so push scored.passes()): LSFT(θ)=−∑i∑t=1Tilog⁡πθ(ai,t∣hi,t).\mathcal{L}_{\text{SFT}}(\theta) = -\sum_i \sum_{t=1}^{T_i} \log \pi_\theta(a_{i,t} \mid h_{i,t}). GRPO (Shao et al. 2024, arXiv:2402.03300). kk = generations rollouts of the same prompt form a group; the advantage is the reward centered on its siblings, and the coefficient is PPO’s clipped surrogate with a KL to the reference: Ai=Ri−mean⁡j∈group(i)Rjstd⁡j∈group(i)Rj,LGRPO=−∑iwi∑tmin⁡ ⁣(ρi,tAi, clip⁡(ρi,t,1−ε,1+ε) Ai)+β KL(πθ ∥ πref).A_i = \frac{R_i - \operatorname{mean}_{j \in \text{group}(i)} R_j}{\operatorname{std}_{j \in \text{group}(i)} R_j}, \qquad \mathcal{L}_{\text{GRPO}} = -\sum_i w_i \sum_{t} \min\!\big(\rho_{i,t} A_i,\ \operatorname{clip}(\rho_{i,t}, 1-\varepsilon, 1+\varepsilon)\, A_i\big) + \beta\, \mathrm{KL}(\pi_\theta \,\|\, \pi_{\text{ref}}). A group whose rollouts all pass or all fail has Ai=0A_i = 0 for every member, which is why select(mode="rl") drops unanimous groups before the GPU is spent. loss_type picks wiw_i, the length weight: grpo averages each sequence’s tokens then averages sequences (wi=1/Tiw_i = 1/T_i); bnpo pools every token in the batch (wi=1/∑jTjw_i = 1/\sum_j T_j); dr_grpo (Liu et al. 2025, arXiv:2503.20783) drops the standard deviation from AiA_i and uses one constant length (wi=1/Lmax⁡w_i = 1/L_{\max}), removing the bias toward long wrong answers. truncated="mask" gives a reply the token cap cut c=0c = 0; "zero" scores it R=0R = 0 instead. Defaults: kk = 8, ε\varepsilon = 0.2, β\beta = 0 (DAPO, Dr. GRPO and CISPO all run without the KL term). DPO (Rafailov et al. 2023, arXiv:2305.18290). A pass ywy_w and a fail yly_l of the same prompt, length matched; no sampling, no reward model: LDPO(θ)=−log⁡σ ⁣(β[log⁡πθ(yw∣x)πref(yw∣x)−log⁡πθ(yl∣x)πref(yl∣x)]),β=0.1.\mathcal{L}_{\text{DPO}}(\theta) = -\log \sigma\!\left(\beta \left[\log \frac{\pi_\theta(y_w \mid x)}{\pi_{\text{ref}}(y_w \mid x)} - \log \frac{\pi_\theta(y_l \mid x)}{\pi_{\text{ref}}(y_l \mid x)}\right]\right), \qquad \beta = 0.1. RM. The same pairs train a scalar head rϕr_\phi under the Bradley-Terry model; reward_model(run) is then a judge: LRM(ϕ)=−log⁡σ ⁣(rϕ(x,yw)−rϕ(x,yl)).\mathcal{L}_{\text{RM}}(\phi) = -\log \sigma\!\big(r_\phi(x, y_w) - r_\phi(x, y_l)\big).

​
Distillation: wai.OPD, wai.OPSD

OPD, on-policy distillation (Agarwal et al. 2023, arXiv:2306.13649; Thinking Machines 2025). The student samples its own reply; a frozen teacher πT\pi_T scores every token of it. The per-token loss is the reverse KL from the student to the teacher over the teacher’s top-kk tokens (Li et al. 2026, arXiv:2604.13016; Fu et al. 2026, arXiv:2603.25562), and its sampled-token form is a policy gradient whose advantage is the log-probability gap: LOPD(θ)=∑i∑t∑v∈top⁡kπT(⋅∣hi,t)πθ(v∣hi,t)log⁡πθ(v∣hi,t)πT(v∣hi,t),ci,t=log⁡πT(ai,t∣hi,t)−log⁡πθ(ai,t∣hi,t).\mathcal{L}_{\text{OPD}}(\theta) = \sum_i \sum_t \sum_{v \in \operatorname{top}_k \pi_T(\cdot \mid h_{i,t})} \pi_\theta(v \mid h_{i,t}) \log \frac{\pi_\theta(v \mid h_{i,t})}{\pi_T(v \mid h_{i,t})}, \qquad c_{i,t} = \log \pi_T(a_{i,t} \mid h_{i,t}) - \log \pi_\theta(a_{i,t} \mid h_{i,t}). No reward, no group, no discount: a token the teacher liked more than the student did is pushed up by the gap. divergence offers forward_kl and jsd for a trainer that has them. Defaults: kk = top_k = 32, samples = 4, temperature 1.0 on both sides so the teacher scores the distribution the student drew from. OPSD, on-policy self-distillation (SDFT, Shenfeld et al. 2026, arXiv:2601.19897; SDPO, Hübotter et al. 2026, arXiv:2601.20802). The same loss, with the teacher the student’s own weights shown a hint pp the student never sees (a passing demonstration, the reference answer, a hint, or a rollout plus the environment’s feedback, privileged): πT(⋅∣hi,t)=πθˉ(⋅∣hi,t⊕pi),θˉ←(1−α) θˉ+α θ(anchor="ema:0.01").\pi_T(\cdot \mid h_{i,t}) = \pi_{\bar\theta}(\cdot \mid h_{i,t} \oplus p_i), \qquad \bar\theta \leftarrow (1-\alpha)\,\bar\theta + \alpha\,\theta \quad (\texttt{anchor="ema:0.01"}). The anchor is what keeps the teacher from chasing the student: "ema:0.01" is the moving average above, "initial" freezes the starting weights (Zhao et al. 2026, arXiv:2601.18734, with forward_kl), "live" is the unanchored student prime-rl runs. samples = 1: there is no group to form.

​
Staleness: wai.Async

Any method above, trained on rollouts that lag the policy by at most off_policy_steps = ss optimizer steps, with one per-token correction for the sampler/trainer gap that exists even at zero lag (Yao et al. 2025): ipo: ci,t←ci,t 1[ ∣πθ(ai,t∣hi,t)−μ(ai,t∣hi,t)∣≤ε ], ε=0.3\texttt{ipo}:\ c_{i,t} \leftarrow c_{i,t}\,\mathbf{1}\big[\,|\pi_\theta(a_{i,t}\mid h_{i,t}) - \mu(a_{i,t}\mid h_{i,t})| \le \varepsilon\,\big],\ \varepsilon = 0.3 icepop: ci,t←ci,t 1[ ρlow≤ρi,t≤ρhigh ], (ρlow,ρhigh)=(0.5,5.0)\texttt{icepop}:\ c_{i,t} \leftarrow c_{i,t}\,\mathbf{1}\big[\,\rho_{\text{low}} \le \rho_{i,t} \le \rho_{\text{high}}\,\big],\ (\rho_{\text{low}}, \rho_{\text{high}}) = (0.5, 5.0) tis: ci,t←ci,t min⁡(ρi,t,ρmax⁡), ρmax⁡=2.0\texttt{tis}:\ c_{i,t} \leftarrow c_{i,t}\,\min(\rho_{i,t}, \rho_{\max}),\ \rho_{\max} = 2.0 ipo masks on probability moved (prime-rl’s default), icepop masks outside a ratio band (Ring-1T, arXiv:2510.18855), tis caps the ratio (verl). Default ss = 8 (ScaleRL, arXiv:2510.13786; one step is free, Noukhovitch et al. 2024, arXiv:2410.18252).

​
The group baselines prime-rl runs: "grpo", "max_rl", "rae"

Three string methods prime_rl_config writes, differing only in AiA_i: grpo: Ai=Ri−mean⁡groupRmax_rl: Ai=Ri−mean⁡groupRmean⁡groupRrae: Ai=Ri−bagent(i),  b←0.95 b+0.05 Ri\texttt{grpo}:\ A_i = R_i - \operatorname{mean}_{\text{group}} R \qquad \texttt{max\_rl}:\ A_i = \frac{R_i - \operatorname{mean}_{\text{group}} R}{\operatorname{mean}_{\text{group}} R} \qquad \texttt{rae}:\ A_i = R_i - b_{\text{agent}(i)},\ \ b \leftarrow 0.95\, b + 0.05\, R_i max_rl (arXiv:2602.02710) divides by the group mean rather than its standard deviation, which makes the gradient unbiased for the maximum-likelihood objective and weights a rarely solved prompt by about 1/p1/p; a group whose mean is 0 carries nothing. rae (SPIRAL, arXiv:2506.24119) is a per-agent running average, the one baseline that stands at one rollout per prompt, and on a single-agent taskset it is REINFORCE against that average.

​
One rollout per prompt: wai.FlashReinforce, wai.SAO, wai.BPCO

The shape a production trace arrives in. With one rollout per prompt every group above is the rollout itself, Ai≡0A_i \equiv 0, and the update is empty. Each of these three replaces the group with something a single trajectory can supply, and each one’s update(batch) returns the cc below as coefficients.

​
FlashReinforce

Baseline from the batch, a gate on drift, equal weight per trajectory. Ai=Ri−1B∑j=1BRj(Eq. 5, no division by a standard deviation)A_i = R_i - \frac{1}{B}\sum_{j=1}^{B} R_j \qquad \text{(Eq. 5, no division by a standard deviation)} di,t=plog⁡pq+(1−p)log⁡1−p1−q,p=μ(ai,t∣hi,t),  q=πθ(ai,t∣hi,t)(Eq. 6)d_{i,t} = p\log\frac{p}{q} + (1-p)\log\frac{1-p}{1-q}, \quad p = \mu(a_{i,t} \mid h_{i,t}),\; q = \pi_\theta(a_{i,t} \mid h_{i,t}) \qquad \text{(Eq. 6)} Dˉi=1Ti∑t=1Tidi,t,mi=1 ⁣[Dˉi≤δ](Eq. 7, 8)\bar D_i = \frac{1}{T_i}\sum_{t=1}^{T_i} d_{i,t}, \qquad m_i = \mathbf{1}\!\left[\bar D_i \le \delta\right] \qquad \text{(Eq. 7, 8)} ci,t=mi Ai ρi,tTi B(the gradient of Eq. 9 with ρ,A,m fixed; Eq. 10)c_{i,t} = \frac{m_i\,A_i\,\rho_{i,t}}{T_i\,B} \qquad \text{(the gradient of Eq. 9 with } \rho, A, m \text{ fixed; Eq. 10)} di,td_{i,t} is the KL between two Bernoullis, “this token versus everything else,” so it needs only the one probability the sampler already stored. A trajectory over the gate drops out whole; the ratio is never clipped, only kept finite. Defaults: δ\delta = trust = 0.003, a lag of off_policy_steps = 8 updates.

​
SAO

A critic in place of the group, credit passed back over tool output, and a band that masks a drifted token instead of clipping it. Over the LL action tokens of one trajectory, in order, with rk=Rr_k = R on the last one and 00 before it, VV the critic’s value at that token and V=0V = 0 past the end: δk=rk+γ V(ak+1)−V(ak),Ak=δk+γλ Ak+1,λ=1−1αL(Eq. 4, 5)\delta_k = r_k + \gamma\,V(a_{k+1}) - V(a_k), \qquad A_k = \delta_k + \gamma\lambda\,A_{k+1}, \qquad \lambda = 1 - \frac{1}{\alpha L} \qquad \text{(Eq. 4, 5)} f(ρ)={ρ1−εlow<ρ<1+εhigh0otherwise(Eq. 2, 3)f(\rho) = \begin{cases} \rho & 1 - \varepsilon_{\text{low}} < \rho < 1 + \varepsilon_{\text{high}} \\ 0 & \text{otherwise} \end{cases} \qquad \text{(Eq. 2, 3)} ci,t=f(ρi,t) Ai,tN,N=action tokens in the batch(Eq. 1)c_{i,t} = \frac{f(\rho_{i,t})\,A_{i,t}}{N}, \qquad N = \text{action tokens in the batch} \qquad \text{(Eq. 1)} The bootstrap V(ak+1)V(a_{k+1}) is the next action token, so an observation between two actions is skipped, not scored. λ\lambda grows with length so the terminal reward reaches the first token with weight λL−1≈e−1/α\lambda^{L-1} \approx e^{-1/\alpha} whatever LL is. No normalization: the critic is the baseline. The critic trains toward the Monte Carlo return, RR at every action token, with a squared-error loss, critic_steps = 2 updates per policy update after critic_warmup = 10 critic-only steps. Defaults: α\alpha = gae_alpha = 1.5, γ\gamma = 1, the band ratio = (1−εlow, 1+εhigh)(1 - \varepsilon_{\text{low}},\, 1 + \varepsilon_{\text{high}}) = (0.7, 6.0).

​
BPCO

A critic that cannot leave the reward’s range, trained on the reward the rollout earned, an advantage left at its natural scale, and a clip on probability rather than on ratio. V=Rmin⁡+(Rmax⁡−Rmin⁡)(12+1πarctan⁡z)(Eq. 9; z the raw head output; bound(z))V = R_{\min} + (R_{\max} - R_{\min})\left(\tfrac{1}{2} + \tfrac{1}{\pi}\arctan z\right) \qquad \text{(Eq. 9; } z \text{ the raw head output; } \texttt{bound(z)}\text{)} V^t=Rfor every token(λV=1, γ=1; Eq. 11)\hat V_t = R \quad \text{for every token} \qquad (\lambda_V = 1,\ \gamma = 1;\ \text{Eq. 11}) δt=rt+V(st+1)−V(st),At=δt+λπAt+1,λπ=1−1αL(Eq. 3, 4, 14; not normalized)\delta_t = r_t + V(s_{t+1}) - V(s_t), \qquad A_t = \delta_t + \lambda_\pi A_{t+1}, \qquad \lambda_\pi = 1 - \frac{1}{\alpha L} \qquad \text{(Eq. 3, 4, 14; not normalized)} surrogatet=min⁡ ⁣(ρtAt, clip⁡ ⁣(ρt, 1−εμt, 1+εμt)At),μt=μ(yt∣st)(Eq. 2, DPPO)\text{surrogate}_t = \min\!\Big(\rho_t A_t,\ \operatorname{clip}\!\big(\rho_t,\ 1 - \tfrac{\varepsilon}{\mu_t},\ 1 + \tfrac{\varepsilon}{\mu_t}\big) A_t\Big), \qquad \mu_t = \mu(y_t \mid s_t) \qquad \text{(Eq. 2, DPPO)} ci,t={ρi,tAi,tthe unclipped branch is the minimum0the clipped branch isc_{i,t} = \begin{cases} \rho_{i,t} A_{i,t} & \text{the unclipped branch is the minimum} \\ 0 & \text{the clipped branch is} \end{cases} The clip range is ε/μt\varepsilon / \mu_t: a token the old policy gave probability 0.05 may move its ratio four times further than one it gave 0.2, because the same ratio change moves less probability mass. The loss sums a sequence’s tokens and averages over sequences, so no length weight enters cc. The critic trains alone for critic_warmup = 15 updates first. Defaults: ε\varepsilon = clip = 0.2, α\alpha = gae_alpha = 0.4, reward_range = (Rmin⁡,Rmax⁡)(R_{\min}, R_{\max}) = (0, 1).

​
From symbol to field

For the three single-rollout methods, update(batch) returns an Update: For the rest, the symbols are the knobs: kk is generations, β\beta is beta, ε\varepsilon is the trainer’s clip, wiw_i is loss_type, ss is off_policy_steps, α\alpha is the rate in anchor, and every one is a named constant in defaults.py with the paper it came from. What the single-rollout three did on a GPU, against their own ablations, is the table in lesson 9 and the three recipes under papers.
Last modified on September 22, 2026