Short Notes on RL

Short Notes on RL for LLMs

Policy Gradients

Why?
When actions are sampled discretely, we generally cannot backpropagate through the action, environment, or reward.
Policy gradients let us optimize the expected reward using only sampled actions, their rewards, and their log-probability gradients.

Key Idea
Sample actions and use the gradients of their log-probabilities. Increase the probability of actions that performed better than the baseline, and decrease the probability of actions that performed worse.

$${\nabla }_{\theta }J(\theta )=\mathbb{E}\left[R{\nabla }_{\theta }\log {\pi }_{\theta }(a\mid s)\right]$$

REINFORCE

Why?
The policy-gradient equation is an expectation. In practice, we only have sampled trajectories. We need a concrete estimator that turns those samples into a gradient update.

Key Idea
Sample a trajectory, observe its return, and use that use the return-to-go $G_t$ to weight the log-probability gradient.

$${\nabla }_{\theta }J(\theta )\approx \sum _t(G_t-b(s_t)){\nabla }_{\theta }\log {\pi }_{\theta }(a_t\mid s_t)$$

KL Divergence

Why?
Because while optimizing a policy, especially an LLM, we often do not want it to drift arbitrarily far from a trusted reference model.

Key Idea
Measure how different one probability distribution is from another.

$$D_{\mathrm{KL}}(\pi \parallel {\pi }_{\mathrm{ref}})={\mathbb{E}}_{a\sim \pi }\left[\log \frac{\pi (a)}{{\pi }_{\mathrm{ref}}(a)}\right]$$

GAE

Why?
REINFORCE’s return-based advantage ($G_t-b(s_t)$) is unbiased but high-variance. A one-step TD (temporal difference) residual is the opposite: low-variance, biased. GAE provides a controllable trade-off between the two.

Key idea
Exponentially-decay-weighted sum of TD residuals, controlled by $\lambda \in [0,1]$:

$${\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)$$ $${\hat{A}}_t^{GAE(\gamma ,\lambda )}=\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+l}$$

• $\lambda =0$ → just ${\delta }_t$ (low variance, high bias)
• $\lambda =1$ → collapses to $G_t-V(s_t)$, the REINFORCE-style advantage (high variance, low bias)

Requires a learned critic $V(s)$, the exact component GRPO later removes.

PPO

Why?
Raw policy-gradient updates can be too large and destabilize training.
The principled fix (a hard trust region, TRPO) needs expensive second-order optimization. PPO wants that stability with plain SGD, reusable across several epochs of the same batch.

Key Idea
Measure how much the new policy has changed relative to the old policy, then prevent excessively large changes.

$${\rho }_t(\theta )=\frac{{\pi }_{\theta }(a_t\mid s_t)}{{\pi }_{{\theta }_{\mathrm{old}}}(a_t\mid s_t)}$$

PPO clips this ratio during optimization.

$$L^{CLIP}(\theta )={\mathbb{E}}_t\left[\min \right({\rho }_t(\theta ){\hat{A}}_t,\text{clip}({\rho }_t(\theta ),1-ϵ,1+ϵ){\hat{A}}_t\left)\right]$$

GAE uses the critic to calculate advantage estimate ${\hat{A}}_t$

$$\text{reward}\to V(s_t)\to \text{GAE}\to {\hat{A}}_t\to \text{PPO update}$$

DPO

Why?
Traditional RLHF with PPO is operationally complicated. It can require a reward model, rollouts, a value model, KL control, and an online RL loop.

Key Idea
Given a preferred response and a rejected response, directly train the model so that the preferred response becomes relatively more likely.

The optimal RLHF policy has a closed form in terms of the reward; plug it into the Bradley–Terry preference model and the reward model cancels out algebraically, leaving a loss on the policy’s own log-probs. No critic, no GAE, no reward model, no rollouts.

$${\mathcal{L}}_{DPO}(\theta )=-{\mathbb{E}}_{(x,y_w,y_l)}\left[\log \sigma \left(\beta \log \frac{{\pi }_{\theta }(y_w\mid x)}{{\pi }_{\text{ref}}(y_w\mid x)}-\beta \log \frac{{\pi }_{\theta }(y_l\mid x)}{{\pi }_{\text{ref}}(y_l\mid x)}\right)\right]$$

GRPO

Why?
PPO usually requires a critic or value model, which is expensive for large language models.

Key Idea
Generate several responses for the same prompt, compare their rewards within that group, and use relative performance as the advantage signal.
Delete the critic and GAE entirely. Sample a group of completions per prompt, score them, and normalize within the group to get the advantage then move KL to the loss instead of the reward.

$${\hat{A}}_i=\frac{R_i-\mathrm{mean}(R_1,\dots ,R_G)}{\mathrm{std}(R_1,\dots ,R_G)}$$ $${\mathcal{J}}_{\mathrm{GRPO}}(\theta )=\mathbb{E}\left[\frac{1}{G}\sum _{i=1}^G\frac{1}{|o_i|}\sum _{t=1}^{|o_i|}\left(\min \left({\rho }_{i,t}(\theta ){\hat{A}}_i,\mathrm{clip}\left({\rho }_{i,t}(\theta ),1-ϵ,1+ϵ\right){\hat{A}}_i\right)-\beta D_{\mathrm{KL}}\left({\pi }_{\theta }\Vert {\pi }_{\mathrm{ref}}\right)\right)\right]$$

Notation

• $G$: number of sampled completions for one prompt; $i$ indexes a completion in the group.
• $o_i$: the $i$-th completion; $|o_i|$ is its token count, and $t$ indexes its tokens.
• ${\hat{A}}_i$: group-normalized reward advantage for completion $i$; the same value is applied to all its tokens.
• $ϵ$: clipping range controlling how far the policy ratio may move from $1$.
• ${\pi }_{\mathrm{ref}}$: fixed reference policy; $D_{\mathrm{KL}}({\pi }_{\theta }\Vert {\pi }_{\mathrm{ref}})$ penalizes drift from it.
• $\beta$: strength of the KL penalty.

References

• Sutton, R.S., McAllester, D., Singh, S., & Mansour, Y. (1999). Policy Gradient Methods for RL with Function Approximation. NeurIPS. https://dblp.org/rec/conf/nips/SuttonMSM99.html
• Williams, R.J. (1992). Simple Statistical Gradient-Following Algorithms for Connectionist RL. Machine Learning, 8, 229–256. https://link.springer.com/article/10.1007/BF00992696
• Kullback, S., & Leibler, R.A. (1951). On Information and Sufficiency. Ann. Math. Statist., 22, 79–86. https://projecteuclid.org/journals/annals-of-mathematical-statistics/volume-22/issue-1/On-Information-and-Sufficiency/10.1214/aoms/1177729694.full
• Schulman, J., Moritz, P., Levine, S., Jordan, M., & Abbeel, P. (2015/2016). High-Dimensional Continuous Control Using Generalized Advantage Estimation. arXiv:1506.02438. https://arxiv.org/abs/1506.02438
• Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347. https://arxiv.org/abs/1707.06347
• Ouyang, L., Wu, J., Jiang, X., et al. (2022). Training Language Models to Follow Instructions with Human Feedback. arXiv:2203.02155. https://arxiv.org/abs/2203.02155
• Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning, C.D., & Finn, C. (2023). Direct Preference Optimization. arXiv:2305.18290. https://arxiv.org/abs/2305.18290
• Shao, Z., Wang, P., Zhu, Q., et al. (2024). DeepSeekMath. arXiv:2402.03300. https://arxiv.org/abs/2402.03300