I feel that when we get used to some form of technology we tend to overlook how fascinating it is that such a technology exists in the first place. Recently I started to wonder how reasoning was added to language models and this led to a rabbit hole that I want you to get into, since I find it truly mind blowing. 

There are already fascinating resources about this. I specially recommend Sebastian Raschka’s Building a Reasoning Model and Nathan Lambert’s RLHF, however, I wanted to write something that compresses both accounts and includes pretty pictures, because I feel this can be better understood in such a way. 

My goal in this post is to introduce the necessary reinforcement learning machinery in order to understand Group Relative Policy Optimization and acquire an intuition for it. Further posts will get into implementation details and eventually running an async RL training loop with real infra. Most of the things covered here are exposed with careful detail in Barto & Sutton classic RL book, but I'm summarizing it for convenience: interested readers will find it rewarding to read the Turing Award winners source material.

With all this in mind, let us start by introducing the very basic definitions of reinforcement learning.

Reinforcement Learning: The Greatest Hits

Throughout this post we will reduce the general reinforcement learning framework to the specific case of LLM's and thus we simplify some things such as the transition function and discretize the state space, but in order to do this we first need to define a Markov Decision Process which is a tuple

where $\mathcal{S}$ contains the possible environment states, $\mathcal{A}$ is the set of actions an agent can take, $\mathcal{T}: \mathcal{S} \times \mathcal{A} \to \mathcal{S}$ is a transition function that defines the environment dynamics, $\mathcal{R}: \mathcal{A} \times \mathcal{S} \to \mathbb{R}$ assigns a reward to each state-action pair and $\gamma \in [0,1]$ is a discount factor.

An agent has no knowledge of $\mathcal{T}$'s definition so it interacts with the environment using a policy $\pi$ which assigns a probability distribution over $\mathcal{A}$ for each $s \in \mathcal{S}$. More precisely

then at each time step $t$ the agent is in state $s$ and samples an action $a_t \sim \pi(\cdot | s_t)$ from its policy. This leads to a reward $r_t = \mathcal{R}(a_t, s_t)$ and the environment transitions to the next state $s_{t+1} = \mathcal{T}(s_t, a_t)$.

Continuing in this fashion an agent travels through a (possibly infinite) trajectory $\tau = (s_0, a_0,r_0, \cdots, s_n, a_n, r_n, \cdots)$ whose reward $\mathcal{R}(\tau)$ is defined as a discounted sum

where $\gamma$ controls the weight of future steps on the final reward.

The ultimate goal of the agent is to find a policy $\pi$ that maximizes the expected cumulative reward

and that's all there is to it (kind of).

RL for LMs

Let us simplify our assumptions a bit. Let $V$ denote some vocabulary (say, tokens) and let $\mathcal{S} = V^*$ i.e. the Kleene star of the vocabulary

This means that now $\mathcal{S}$ contains all possible token sequences. Making $\mathcal{A} = V$ we can turn $\mathcal{T}$ into a deterministic function by defining it as concatenation, for this denote $s \in \mathcal{S}$ by $s = (v_1, \cdots, v_n)$, then

Thus at each step the agent is picking a token and appending it to the state.

Suppose $\mathcal{R}$ rewards trajectories that end up with compiling code or maybe with an appropriate proof of a theorem. Since we want to grade the entire chain of thought we also set $\gamma = 1$.

For language models the policy $\pi$ is the model itself because it is assigning probabilities to each next token and since each language model is characterized by its weights $\theta$ we can parameterize policies by them, so we usually write $\pi_\theta$ to represent a language model with weights $\theta$.

In this setting, the objective is to find the weights that maximize the expected cumulative reward

and here it gets tricky, but tricky is fun.

What is a Policy Gradient?

A detailed account of the policy gradient derivation can be found in chapter six of the RLHF book, but the short version is that we can frame our goal as maximizing the following objective function

a professor of mine once told me: if you are doing differential geometry and you don't know what to do in order to solve a problem, try taking a derivative. This is very wise advice so our instincts are screaming at us reach for the gratient $\nabla_\theta J(\theta)$! but what does this expression look like? I'll spare you the computation (really, go read the book), the result ends up being

Each term is telling us:

• $\mathcal{R}(\tau)$: how good was the trajectory?
• $\nabla_\theta \log \pi_\theta(a_t | s_t)$: what direction in parameter space makes action $a_t$ more likely?

Moving in such a direction makes good trajectories more likely to happen and bad trajectories unlikely. Very cool, but then this turns into the neverending problem of finding a way to compute it efficiently.

What is Group Relative Policy Optimization?

It seems to be the case that we are bound by time and so are computers, meaning that we can't sample all possible trajectories of a policy $\pi_\theta$. Hence we must find a way to get just enough information about this trajectory space so that our $J(\theta)$ objective function is well approximated.

The Policy Optimization part is clear, so let's focus on the Group Relative part of GRPO. The idea is very simple, we sample $G$ trajectories (which we also call rollouts) and then we compute how each one is performing compared to the rest i.e. relative to the group, we then weight the log probabilities of each action by this advantage and do the gradient over this measure.

Let $r_i$ be the reward of completing the trajectory $\tau_i$ then the advantage is set to be

and we can use this advantage to incentivize the generation of good trajectories. Now, as above, we do this by multiplying it with the log probabilities for each trajectory and end up with an objective function

If you go an read the original GRPO paper you will see that the objective function there is similar to ours but several terms are missing, this is because we are just giving a mathematical framework, further considerations such as the cost of each rollout and stale policies will be explored in the next post.