This repo serves as a workspace for my learning of RL in the context of modern LLMs.

Resources:

• https://spinningup.openai.com/
• https://rlhfbook.com/c/11-policy-gradients.html

Model-based: Agent learns to model the environment. This allows the agent to plan ahead, but the learned model might not generalise well.

Policy gradient is the gradient of Expected return $J(\pi_\theta)$ (our objective) with respect to Policy's parameters $\theta$. We can derive a nice form for it.

Note: A trajectory's log-prob is equal to log-prob of initial state, plus the log-prob sum of (action + state transition). Since initial state and state transition are properties of the environment, they do not depend on policy's parameters. Therefore, the derivative of trajectory's log-prob (wrt policy's parameters) is equal to the sum of log-policy's derivatives.

Note that the total number of timesteps $T$ needs not to be the same across trajectories.

In the sampling form (estimator)

Notice that once we sample the trajectories (using current policy), we can bring the gradient operator out of the sum. We can define the objective function:

Please note that this loss function is only defined to simplify the computation of policy gradient using an autograd engine. THERE IS NO MEANING behind this loss function. Do not try to interpret its value.

Other forms of Policy gradient It turns out that we can replace $G(\tau)$ in Policy gradient equation with other terms that result in the same expected value. These alternatives can offer a tighter variance of Policy gradient estimator, making training more data-efficient. Here are some alternative forms (proofs won't be covered).

• Using on-policy Q-function

• Using Advantage function - this is the most common since there are many ways to estimate the advantage function and it results in lowest variance.

Notice in the above two forms, all the terms inside the summation depends only on current $(s_t,a_t)$. Hence, we can specify a loss depending only on those.

One common way to estimate the Value function is with a neural network. We will train the Value function network together with the Policy network, using the same rollout data.

Thanks to this, we can estimate the Value function at current state without integrating all possible trajectories.

Generalized Advantage Estimator (GAE) https://arxiv.org/abs/1506.02438

Note that the Value function above is estimated as well, using least mean squares optimization mentioned previously.

When $\lambda=0$

This expression is TD-residual for the estimated Value function. $r_t + \gamma V(s_{t+1})$ is basically the Q-function if the $V(s_t)$ is a perfect estimate. This Advantage estimator has high bias since the Value function estimator is not accurate.

When $\lambda=1$

The summation term is Return of a particular trajectory sample starting at current state and action. This is an unbiased estimator of Q-function, but it has high variance, which we are trying to avoid.

https://arxiv.org/abs/1502.05477

Impose an additional constraint not to deviate too much from the old policy -> avoid big update that can collapse training.

Surrogate advantage is a measure of how policy $\pi_\theta$ performs relative to the old policy $\pi_{\theta_{old}}$, using data from the old policy.

Average KL-divergence over states visited by the old policy

Theoretical TRPO update

Note

• Gradient of surrogate advantage function (wrt policy's parameters) is still equal to policy gradient.
• Surrogate advantage is taking expectation over $(s,a)\sim\pi_{\theta_{old}}$, instead of summing over the time steps then taking expectation over trajectories. Not sure if the author is assuming $T$ is the same across different trajectories. TODO: double-check this

https://arxiv.org/abs/1707.06347

PPO-clip Loss function

Note that since Advantage $A^{\pi_{\theta_{old}}}$ can be either positive or negative, we can't factor out the Advantage term.

PPO-Penalty (not used much) Loss function

Where the penalty coefficient $\beta$ is adjusted automatically over the course of training.

Reinforcement Learning from Human Feedback (RLHF) + PPO (OpenAI)

• https://arxiv.org/abs/1909.08593 / https://github.com/openai/lm-human-preferences
• https://arxiv.org/abs/2009.01325 / https://github.com/openai/summarize-from-feedback
• https://arxiv.org/abs/2203.02155 (InstructGPT)

Given prompt $x$ and completion $y$. We consider this to be a contextual bandit environment i.e. single-step RL, since (1) user prompt $x$ is the context, and (2) we only have reward on the full response. Action space is all possible $y\in V^{|y|}$, where $V$ is the vocabulary of the LLM.

The Reward function is set as

Where $r_\phi(x,y)$ is the Reward model, and the KL divergence term is used to keep the RL model not deviate too much from the SFT model.

The Reward model is frozen during RL fine-tuning. InstructGPT uses the following objective (very similar to Discriminator loss in Relativistic GAN), where $y_1$ is preferred over $y_0$. TODO: check how data is collected

The OpenAI papers above then use PPO-Clip to RL-finetune the model with this engineered reward. GAE is used to estimate the Advantage function, which requires estimating the Value function with a neural network. This Value function is initialized from the Reward model.

TODO: check OpenAI code to see how they normalize the loss across tokens / samples.

Direct Preference Optimization (DPO) https://arxiv.org/abs/2305.18290

Use Reward model's objective to optimize the LLM directly

Reinforcement Learning with Verifiable Rewards (RLVR) First coined in Tulu3 (https://arxiv.org/abs/2411.15124). The reward function comes from an easily verifiable procedure e.g. answer to a math problem.

Group Relative Policy Optimization (GRPO) https://arxiv.org/abs/2402.03300

In this setup, state is prefix (prompt $x$ + tokens generated so far $y_{1..t-1}$), and action is next token $y_t$. Number of timesteps is equal to number of generated tokens.

For a given prompt $x$, sample a few responses ${y}$ using the old policy. GRPO loss defined on a single $(x,y)$ pair is:

This is basically PPO-Clip loss with extra KL penalty term (the KL term is opposite direction of PPO-Penalty). Advantage is estimated based on relative rewards within a group (hence the name). Also notice that KL penalty is not part of the reward, like RLHF, but part of the loss. TODO: check how KL term is computed.

If we only have reward at the end of each rollout $r(x,y)$, Advantage is estimated as normalized rewards.

If we have rewards at intermediate steps (process supervision), Advantage is estimated as the sum of current and future normalized rewards. TODO: check how reward normalization is done in this case. Each rollout can have different number of (intermediate) rewards.

The full GRPO loss is then averaged over sampled responses $y\sim\pi_{\theta_{old}}$ for a single $x$, then averaged over sampled prompts $x$.

TODO: check DeepSeek-R1 as well

GRPO Done Right (Dr. GRPO) https://arxiv.org/abs/2503.20783

Length normalization is removed

Reward normalization only shifts mean to zero. Don't rescale it to unit variance

Decoupled Clip and Dynamic Sampling Policy Optimization (DAPO) https://arxiv.org/abs/2503.14476

On a prompt-response pair $(x,y)$, DAPO modifies PPO loss to have 2 separate clipping values

Similar to GRPO, Advantage is estimated as normalized rewards. The other important difference is

1. How sampling is done: Remove all prompts that have 0 or 1 accuracy i.e. sampled responses are all wrong or all correct. This is to avoid normalized reward being zero, which will contribute to training instability.
2. How expectation is computed: Loss is averaged over all tokens in a batch, instead of:

• GRPO: loss for each response is averged over its length (per-sample), then averaged over rollouts.
• Dr. GRPO: loss for each response is summed over its length (per-sample), then averaged over rollouts.

3. Special treatment for truncated responses: Long responses are truncated. DAPO will either ignore truncated samples (overlong filtering) or apply a penalty to the reward function (soft overlong punishment).

Unlike GRPO and OpenAI RLHF, DAPO doesn't include any KL penalty term at all, whether in its reward function or loss function.