Deriving the PPO Loss from First Principles

I have been trying to wrap my head around reinforcement learning methods like DPO, GRPO, and RLVR for a while now, especially with all the recent work showing how effective they can be for LLM post-training. Since I amm still pretty new to RL, I figured the best place to start was Proximal Policy Optimization (PPO), the algorithm OpenAI used to show how reinforcement learning could meaningfully improve LLM alignment (InstructGPT paper). My hope is that getting comfortable with PPO will give me the right mental model for the policy-gradient side of things and make it easier to understand the newer LLM-specific RL methods built on similar ideas.

If you start learning RL, you quickly realize it involves a lot of math! So I decided to lean into that and do a few (possibly annoying) derivation sessions to really understand the PPO objective by building it up from first principles, similar to how Umar Jamil does in his video.

A huge shoutout to Umar Jamil's video on RLHF and PPO: it was incredibly helpful for building intuition and understanding the math behind the PPO loss.

Below is my attempt at the derivation based on the original PPO and InstructGPT papers and Umar Jamil’s video.

I: Reinforcement Learning: Core Definitions

II: Reward Model in RLHF for LLMs

A Reward Model (RM) is a neural network that takes a prompt $xx$ and a response $yy$ as input and outputs a scalar reward $r_{\varphi }(x,y)\in \mathbb{R}r\_\backslash phi(x,\; y)\; \backslash in\; \backslash mathbb\{R\}$ indicating how "good" or "aligned" that response is according to human preferences.

Policy-gradient methods (including PPO) require a scalar objective to update the policy parameters. In standard RL, the environment provides this signal. However for language generation, there is no natural environment giving us rewards for "good" responses. Having humans rate every output is impractical and for gradient-based optimization, we need a differentiable scalar signal to backpropagate through. Thus, we require a cheap, differentiable proxy for human preferences during RL training. A learned RM provides exactly this.

How is the Reward Model Trained?

The standard procedure for training the reward model is:

1. Sample prompts ( $xx$ )
2. Generate multiple candidate completions ( $y_1,y_2,\dots ,y_{K}y\_1,\; y\_2,\; \backslash ldots,\; y\_K$ ) from a baseline policy (often an SFT model).
3. Ask humans to compare candidates (pairwise preferences are easier than absolute scoring).
4. Train the RM ( $r_{\varphi }r\_\backslash phi$ ) to predict those preferences.

Architecturally, the reward model is typically:

• Initialized from a pretrained language model (often the SFT model itself)
• The final non-embedding layer (which projects to vocabulary) is removed
• Replaced it with a linear layer that projects the hidden state of the last token to a single scalar output

Reward Model Loss Function

The reward model is trained using the Bradley-Terry model for pairwise comparisons. The probability that response $y_{w}y\_w$ (preferred) is preferred over $y_{l}y\_l$ (less preferred) for any prompt $xx$ is modeled as:

$$\begin{array}{cccc}& P(y_w\succ y_l\mathrm{\mid }x)=\sigma (r_{\theta }(x,y_w)-r_{\theta }(x,y_l))& & \text{(II.I)}\end{array}P(y\_w\; \backslash succ\; y\_l\; |\; x)\; =\; \backslash sigma\backslash left(r\_\backslash theta(x,\; y\_w)\; -\; r\_\backslash theta(x,\; y\_l)\backslash right)\; \backslash tag\{II.I\}$$

where $\sigma \backslash sigma$ is the sigmoid function: $\sigma (z)=\frac{1}{1+e^{-z}}\backslash sigma(z)\; =\; \backslash frac\{1\}\{1\; +\; e^\{-z\}\}$

The negative log-likelihood loss is:

$${\mathcal{L}}_{\text{RM}}(\varphi )=-{\mathbb{E}}_{(x,y_w,y_l)\sim \mathcal{D}}[\log \sigma (r_{\varphi }(x,y_w)-r_{\varphi }(x,y_l))]\backslash mathcal\{L\}\_\{\backslash text\{RM\}\}(\backslash phi)\; =\; -\backslash mathbb\{E\}\_\{(x,\; y\_w,\; y\_l)\; \backslash sim\; \backslash mathcal\{D\}\}\; \backslash left[\; \backslash log\; \backslash sigma\backslash left(r\_\backslash phi(x,\; y\_w)\; -\; r\_\backslash phi(x,\; y\_l)\backslash right)\; \backslash right]$$

One can verify that this loss forces the reward model to assign higher rewards to preferred responses (see InstructGPT paper or Umar Jamil's video for a detailed walkthrough).

There are two key insights here:

1. We don't need absolute scores, we only need the reward model to correctly rank responses.
2. The loss depends only on differences ( $r_{\varphi }(x,y_w)-r_{\varphi }(x,y_l)r\_\backslash phi(x,\; y\_w)\; -\; r\_\backslash phi(x,\; y\_l)$ ), so it is invariant to adding a constant to all rewards. This will be useful later when we discuss the PPO loss.

The reward model serves as a learned proxy for human preferences, converting the intractable problem of getting human feedback on every generation into a tractable supervised learning problem. Once trained, it provides the scalar signal $r_{\varphi }(x,y)r\_\backslash phi(x,\; y)$ needed to optimize our policy (LLM) using rl algorithms like PPO.

III: Trajectories and Returns

Trajectory

A trajectory (also called a rollout or episode) is a sequence of states ( $ss$ ), actions ( $aa$ ), and rewards ( $rr$ ) generated by an agent interacting with an environment:

$$\tau =(s_0,a_0,r_0,s_1,a_1,r_1,\dots ,s_T,a_T,r_T)\backslash tau\; =\; (s\_0,\; a\_0,\; r\_0,\; s\_1,\; a\_1,\; r\_1,\; \backslash ldots,\; s\_T,\; a\_T,\; r\_T)$$

In the context of LLMs, a trajectory corresponds to the entire sequence of token generations. It is the prompt followed by all generated tokens until the end-of-sequence token.

Note that the states are always stochastically modeled, and $s_{t+1}s\_\{t+1\}$ can be represented as $s_{t+1}\sim P(s_{t+1}\mathrm{\mid }{s}_t,a_t)s\_\{t+1\}\; \backslash sim\; P(s\_\{t+1\}\; |\; s\_t,\; a\_t)$. Given a stochastic policy ${\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\backslash pi\_\backslash theta\; (a\_t\; |\; s\_t)$, the probability of a trajectory $\tau \backslash tau$ is the product of:

1. The initial state distribution ${\rho }_0(s_0)\backslash rho\_0(s\_0)$
2. The stochastic policy ${\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\backslash pi\_\backslash theta\; (a\_t\; |\; s\_t)$
3. The environment transition dynamics $P(s_{t+1}\mathrm{\mid }{s}_t,a_t)P(s\_\{t+1\}\; |\; s\_t,\; a\_t)$

$$\begin{array}{cccc}& P(\tau \mathrm{\mid }{\pi }_{\theta })={\rho }_0(s_0)\prod _{t=0}^{T-1}{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot P(s_{t+1}\mathrm{\mid }{s}_t,a_t)& & \text{(III.I)}\end{array}P(\backslash tau\; |\; \backslash pi\_\backslash theta)\; =\; \backslash rho\_0(s\_0)\; \backslash prod\_\{t=0\}^\{T-1\}\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; P(s\_\{t+1\}\; |\; s\_t,\; a\_t)\; \backslash tag\{III.I\}$$

Return

The return is the cumulative reward collected over the full trajectory ( $\tau \backslash tau$ ). The simplest form is the undiscounted return:

$$R(\tau )=\sum _{t=0}^T{r}_{t}R(\backslash tau)\; =\; \backslash sum\_\{t=0\}^\{T\}\; r\_t$$

More generally, we use the discounted return:

$$\begin{array}{cccc}& R(\tau )=\sum _{k=0}^{\mathrm{\infty }}{\gamma }^k{r}_k=r_0+\gamma r_1+{\gamma }^2{r}_2+\cdots & & \text{(III.II)}\end{array}R(\backslash tau)\; =\; \backslash sum\_\{k=0\}^\{\backslash infty\}\; \backslash gamma^k\; r\_\{k\}\; =\; r\_0\; +\; \backslash gamma\; r\_\{1\}\; +\; \backslash gamma^2\; r\_\{2\}\; +\; \backslash cdots\; \backslash tag\{III.II\}$$

where $\gamma \in [0,1]\backslash gamma\; \backslash in\; [0,\; 1]$ is the discount factor. The discount factor $\gamma \backslash gamma$ serves a couple of purposes:

1. It ensures the return is finite for infinite-horizon tasks ( $T\to \mathrm{\infty }T\backslash to\backslash infty$ ).
2. It prioritizes immediate rewards over distant ones.

IV: Policy Gradient Optimization and REINFORCE Algorithm

The goal of reinforcement learning is to find a policy ${\pi }_{\theta }\backslash pi\_\backslash theta$ that maximizes the expected return over all possible trajectories:

$$\begin{array}{cccc}& \boxed{{\displaystyle J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]}}& & \text{(IV.I)}\end{array}\backslash boxed\{J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}[R(\backslash tau)]\}\; \backslash tag\{IV.I\}$$

This is our objective function and we want to find parameters ${\theta }^{\ast }\backslash theta^*$ such that:

$${\theta }^{\ast }=\arg \underset{\theta }{\max }J(\theta )\backslash theta^*\; =\; \backslash arg\backslash max\_\backslash theta\; J(\backslash theta)$$

To maximize $J(\theta )J(\backslash theta)$ using gradient-based methods, we need to compute ${\mathrm{\nabla }}_{\theta }J(\theta )\backslash nabla\_\backslash theta\; J(\backslash theta)$ and perform gradient ascent:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\theta }_{k+1}={\theta }_k+\alpha {{\mathrm{\nabla }}_{\theta }J({\pi }_{\theta })\mid }_{{\theta }_k}}}& & \text{(IV.II)}\end{array}\backslash boxed\{\backslash theta\_\{k+1\}\; =\; \backslash theta\_k\; +\; \backslash alpha\; \backslash left.\; \backslash nabla\_\backslash theta\; J(\backslash pi\_\backslash theta)\; \backslash right|\_\{\backslash theta\_k\}\}\; \backslash tag\{IV.II\}$$

This policy gradient looks simple in equation form but it is intractable to compute. The expectation is over trajectories sampled from ${\pi }_{\theta }\backslash pi\_\backslash theta$, which itself depends on $\theta \backslash theta$. We can't simply enumerate all possible trajectories. This is computationally intractable for any reasonably sized state-action space (and certainly not possible for LLMs!).

Thus, as a next step we need to derive some sort of reasonable and tractable approximation for ${\mathrm{\nabla }}_{\theta }J(\theta )\backslash nabla\_\backslash theta\; J(\backslash theta)$. We do this by using the log-derivative trick.

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathrm{\nabla }}_{\theta }{\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[R(\tau )]\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash nabla\_\backslash theta\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}[R(\backslash tau)]$$

This expectation can be written as an integral: $$={\mathrm{\nabla }}_{\theta }{\int }_{\tau }P(\tau \mathrm{\mid }\theta )R(\tau )d\tau =\; \backslash nabla\_\backslash theta\; \backslash int\_\backslash tau\; P(\backslash tau\; |\; \backslash theta)\; R(\backslash tau)\; \backslash ,\; d\backslash tau$$

Bringing the gradient inside the integral: $$={\int }_{\tau }{\mathrm{\nabla }}_{\theta }P(\tau \mathrm{\mid }\theta )R(\tau )d\tau =\; \backslash int\_\backslash tau\; \backslash nabla\_\backslash theta\; P(\backslash tau\; |\; \backslash theta)\; R(\backslash tau)\; \backslash ,\; d\backslash tau$$

Now we apply the log-derivative trick: $${\mathrm{\nabla }}_{\theta }\log P(\tau \mathrm{\mid }\theta )=\frac{{\mathrm{\nabla }}_{\theta }P(\tau \mathrm{\mid }\theta )}{P(\tau \mathrm{\mid }\theta )}\backslash nabla\_\backslash theta\; \backslash log\; P(\backslash tau\; |\; \backslash theta)\; =\; \backslash frac\{\backslash nabla\_\backslash theta\; P(\backslash tau\; |\; \backslash theta)\}\{P(\backslash tau\; |\; \backslash theta)\}$$

Rearranging: ${\mathrm{\nabla }}_{\theta }P(\tau \mathrm{\mid }\theta )=P(\tau \mathrm{\mid }\theta ){\mathrm{\nabla }}_{\theta }\log P(\tau \mathrm{\mid }\theta )\backslash nabla\_\backslash theta\; P(\backslash tau\; |\; \backslash theta)\; =\; P(\backslash tau\; |\; \backslash theta)\; \backslash nabla\_\backslash theta\; \backslash log\; P(\backslash tau\; |\; \backslash theta)$ and substituting back, we get: $$={\int }_{\tau }P(\tau \mathrm{\mid }\theta ){\mathrm{\nabla }}_{\theta }\log P(\tau \mathrm{\mid }\theta )R(\tau )d\tau =\; \backslash int\_\backslash tau\; P(\backslash tau\; |\; \backslash theta)\; \backslash nabla\_\backslash theta\; \backslash log\; P(\backslash tau\; |\; \backslash theta)\; R(\backslash tau)\; \backslash ,\; d\backslash tau$$

which can also be written as the following expectation:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[{\mathrm{\nabla }}_{\theta }\log P(\tau \mathrm{\mid }\theta )\cdot R(\tau )]}}& & \text{(IV.III)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\; \backslash nabla\_\backslash theta\; \backslash log\; P(\backslash tau\; |\; \backslash theta)\; \backslash cdot\; R(\backslash tau)\; \backslash right]\}\; \backslash tag\{IV.III\}$$

Note, here the gradient is now the expectation of the gradient of the log-probability of the trajectory. This can further be simplified by using the trajectory probability expression (III.I): $$P(\tau \mathrm{\mid }\theta )={\rho }_0(s_0)\prod _{t=0}^{T-1}{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot P(s_{t+1}\mathrm{\mid }{s}_t,a_t)P(\backslash tau\; |\; \backslash theta)\; =\; \backslash rho\_0(s\_0)\; \backslash prod\_\{t=0\}^\{T-1\}\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; P(s\_\{t+1\}\; |\; s\_t,\; a\_t)$$

Taking the log:

$$\log P(\tau \mathrm{\mid }\theta )=\log {\rho }_0(s_0)+\sum _{t=0}^{T-1}\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)+\sum _{t=0}^{T-1}\log P(s_{t+1}\mathrm{\mid }{s}_t,a_t)\backslash log\; P(\backslash tau\; |\; \backslash theta)\; =\; \backslash log\; \backslash rho\_0(s\_0)\; +\; \backslash sum\_\{t=0\}^\{T-1\}\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; +\; \backslash sum\_\{t=0\}^\{T-1\}\; \backslash log\; P(s\_\{t+1\}\; |\; s\_t,\; a\_t)$$

When we take ${\mathrm{\nabla }}_{\theta }\backslash nabla\_\backslash theta$, only the policy term depends on $\theta \backslash theta$:

$${\mathrm{\nabla }}_{\theta }\log P(\tau \mathrm{\mid }\theta )=\sum _{t=0}^{T-1}{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\backslash nabla\_\backslash theta\; \backslash log\; P(\backslash tau\; |\; \backslash theta)\; =\; \backslash sum\_\{t=0\}^\{T-1\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)$$

The initial state distribution and transition dynamics are independent of $\theta \backslash theta$, so their gradients vanish. Substituting back, we obtain the policy gradient theorem:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot R(\tau )]}}& & \text{(IV.IV)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\; \backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; R(\backslash tau)\; \backslash right]\}\; \backslash tag\{IV.IV\}$$

This is a remarkable result. We can compute the gradient of our objective without differentiating through the environment dynamics and only need gradients of the log-probabilities of our policy.

Since we cannot compute the expectation exactly, we approximate it with a sample mean by sampling $NN$ trajectories:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )\approx \widehat{g}=\frac{1}{N}\sum _{i=1}^N(\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_{i,t}\mathrm{\mid }{s}_{i,t}))R({\tau }_i)}}& & \text{(IV.V)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; \backslash approx\; \backslash hat\{g\}\; =\; \backslash frac\{1\}\{N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash left(\; \backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_\{i,t\}\; |\; s\_\{i,t\})\; \backslash right)\; R(\backslash tau\_i)\}\; \backslash tag\{IV.V\}$$

This gives us the REINFORCE algorithm:

1. Initialize: Start with a pretrained or supervised fine-tuned (SFT) language model ${\pi }_{\theta }\backslash pi\_\backslash theta$

2. Sample prompts: Draw a batch of $NN$ prompts $\{x_1,x_2,\dots ,x_N\}\backslash \{x\_1,\; x\_2,\; \backslash ldots,\; x\_N\backslash \}$ from a dataset

3. Generate trajectories: For each prompt $x_{i}x\_i$, generate a response $y_i=(a_0,a_1,\dots ,a_T)y\_i\; =\; (a\_0,\; a\_1,\; \backslash ldots,\; a\_T)$ by sampling tokens from the policy ${\pi }_{\theta }\backslash pi\_\backslash theta$. Each trajectory is the sequence of states (prompt + generated tokens so far) and actions (selected tokens).

4. Compute log-probabilities: For each trajectory, compute the log-probability of each generated token given its context:

$$\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\text{for }t=0,1,\dots ,T\backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash quad\; \backslash text\{for\; \}\; t\; =\; 0,\; 1,\; \backslash ldots,\; T$$

6. Compute rewards: Score each complete (prompt, response) pair using the reward model: $R({\tau }_i)=r_{\varphi }(x_i,y_i)R(\backslash tau\_i)\; =\; r\_\backslash phi(x\_i,\; y\_i)$

7. Estimate policy gradient: Compute the gradient estimate using (IV.V): $$\widehat{g}=\frac{1}{N}\sum _{i=1}^N(\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_{i,t}\mathrm{\mid }{s}_{i,t}))R({\tau }_i)\backslash hat\{g\}\; =\; \backslash frac\{1\}\{N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash left(\; \backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_\{i,t\}\; |\; s\_\{i,t\})\; \backslash right)\; R(\backslash tau\_i)$$

8. Update policy: Perform a gradient ascent step: $\theta \leftarrow \theta +\alpha \widehat{g}\backslash theta\; \backslash leftarrow\; \backslash theta\; +\; \backslash alpha\; \backslash hat\{g\}$

9. Repeat: Go back to Step 2 and iterate until convergence

While REINFORCE provides an unbiased gradient estimate, it suffers from two critical issues that make it impractical for LLM training:

1. High Variance: The gradient estimate $\widehat{g}\backslash hat\{g\}$ suffers from high variance depending on the sampled trajectories. This variance can be large and can lead to noisy gradients and unstable training.

   If you look again at (IV.V), the gradient estimate for each action is weighted by the return of the entire trajectory $R(\tau )R(\backslash tau)$. This means that even if an action was good, it might receive a negative gradient update simply because other actions in the trajectory led to poor outcomes (or vice versa). Over many samples, the noise introduced by this coupling can be substantial, leading to high variance

2. On-Policy Constraint (Sample Inefficiency): REINFORCE requires trajectories sampled from the current policy ${\pi }_{\theta }\backslash pi\_\backslash theta$. Thus after every gradient update, previously collected trajectories must be discarded and new ones need to be sampled from the updated policy. For LLMs, where each trajectory requires a full forward pass through a billion(s)-parameter model, this is prohibitively expensive especially when we need many small gradient steps to train effectively.

V: Reducing Variance and the Advantage Function

The REINFORCE algorithm provides an unbiased gradient estimate (IV.V). However while unbiased, this estimator suffers from high variance.

Replacing Full-Trajectory Return with Reward-to-Go (using causality)

A first variance reduction comes from noticing that action $a_{t}a\_t$ taken at time $tt$ cannot influence rewards that were received before time $tt$. This is a fundamental consequence of causality. These past reward terms contribute only noise to the gradient estimate and add variance without contributing any signal. Thus, we can remove them and consider only the rewards-to-go :

$${\widehat{R}}_t=\sum _{t^{\mathrm{\prime }}=t}^T{r}_{t^{\mathrm{\prime }}}\backslash hat\{R\}\_t\; =\; \backslash sum\_\{t\text{'}=t\}^\{T\}\; r\_\{t\text{'}\}$$

This gives us a lower-variance estimator:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_{i,t}\mathrm{\mid }{s}_{i,t})\cdot {\widehat{R}}_{i,t}}}& & \text{(V.I)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; \backslash approx\; \backslash frac\{1\}\{N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_\{i,t\}\; |\; s\_\{i,t\})\; \backslash cdot\; \backslash hat\{R\}\_\{i,t\}\}\; \backslash tag\{V.I\}$$

where ${\widehat{R}}_{i,t}={\sum }_{t^{\mathrm{\prime }}=t}^T{r}_{i,t^{\mathrm{\prime }}}\backslash hat\{R\}\_\{i,t\}\; =\; \backslash sum\_\{t\text{'}=t\}^\{T\}\; r\_\{i,t\text{'}\}$ is the rewards-to-go for trajectory $ii$ starting from time $tt$.

Subtracting a Baseline

A second complementary technique for variance reduction is to subtract a baseline $b(s_t)b(s\_t)$ from the rewards. The key insight is that we can subtract any function that does not depend on the action from our reward signal without changing the expected value of the gradient.

Thus we can subtract a state-dependent baseline $b(s_t)b(s\_t)$ from our rewards-to-go to yield an unbiased gradient estimator:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_{i,t}\mathrm{\mid }{s}_{i,t})\cdot ({\widehat{R}}_{i,t}-b(s_{i,t}))}}& & \text{(V.II)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; \backslash approx\; \backslash frac\{1\}\{N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_\{i,t\}\; |\; s\_\{i,t\})\; \backslash cdot\; \backslash left(\backslash hat\{R\}\_\{i,t\}\; -\; b(s\_\{i,t\})\backslash right)\}\; \backslash tag\{V.II\}$$

Value Functions: $V^{\pi }(s)V^\backslash pi(s)$ and $Q^{\pi }(s,a)Q^\backslash pi(s,\; a)$

The baseline is still an arbitrary function. To make it more systematic and concrete, there are two fundamental functions from RL theory.

State Value Function: The state value function $V^{\pi }(s)V^\backslash pi(s)$ is the expected return when the agent is in state $ss$ and acts according to policy $\pi \backslash pi$:

$$V^{\pi }(s)={\mathbb{E}}_{\tau \sim \pi }\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t\right|s_0=s]V^\backslash pi(s)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\}\backslash left[\backslash sum\_\{t=0\}^\{\backslash infty\}\; \backslash gamma^t\; r\_t\; \backslash ;\backslash middle|\backslash ;\; s\_0\; =\; s\backslash right]$$

Intuitively, $V^{\pi }(s)V^\backslash pi(s)$ tells "How good is this state on average?" and is used as a baseline $b(s)=V^{\pi }(s)b(s)\; =\; V^\backslash pi(s)$.

Action Value Function (Q-function): The action value function $Q^{\pi }(s,a)Q^\backslash pi(s,\; a)$ is the expected return when starting in state $ss$ and taking action $aa$ and then acting according to policy $\pi \backslash pi$:

$$Q^{\pi }(s,a)={\mathbb{E}}_{\tau \sim \pi }\left[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{r}_t\right|s_0=s,a_0=a]Q^\backslash pi(s,\; a)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\}\backslash left[\backslash sum\_\{t=0\}^\{\backslash infty\}\; \backslash gamma^t\; r\_t\; \backslash ;\backslash middle|\backslash ;\; s\_0\; =\; s,\; a\_0\; =\; a\backslash right]$$

Intuitively, $Q^{\pi }(s,a)Q^\backslash pi(s,\; a)$ tells "How good is this specific action in this state?" and in RL, the rewards-to-go is estimated as $Q^{\pi }(s,a)Q^\backslash pi(s,\; a)$.

In the LLM context:

• $V^{\pi }(s)V^\backslash pi(s)$ estimates the expected reward for a given prompt + partial response, assuming the model continues generating according to its current policy.
• $Q^{\pi }(s,a)Q^\backslash pi(s,\; a)$ estimates the expected reward if, from the current prompt + partial response, the model generates a specific next token $aa$ and then continues according to its policy.

Advantage Function

The advantage function $A^{\pi }(s,a)A^\backslash pi(s,\; a)$ measures how much better (or worse) a specific action $aa$ is compared to the average action under the policy:

$$\begin{array}{cccc}& \boxed{{\displaystyle A^{\pi }(s,a)=Q^{\pi }(s,a)-V^{\pi }(s)}}& & \text{(V.III)}\end{array}\backslash boxed\{A^\backslash pi(s,\; a)\; =\; Q^\backslash pi(s,\; a)\; -\; V^\backslash pi(s)\}\; \backslash tag\{V.III\}$$

The advantage function directly tells us: "How much better is this particular action compared to what we would typically do in this state?" This is precisely the signal we want for policy improvement. We want to increase the probability of actions with positive advantage and decrease the probability of actions with negative advantage.

From Umar Jamil's video:
In the LLM context consider a state where the prompt is "Where is Shanghai?" and the model has generated "Shanghai is". From this state:

• If the model samples the token "in" (leading toward "Shanghai is in China"), this action likely has positive advantage. This is because it is better than the average token the model might produce.
• If the model samples the token "delicious" (leading toward an incoherent response), this action likely has negative advantage. This is because it is worse than the average token the model might produce.

Advantage-Weighted Policy Gradient

Substituting the rewards-to-go and the value function as a baseline, we get the following form of the policy gradient: $${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot (Q^{\pi }(s_t,a_t)-V^{\pi }(s_t))]\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; (Q^\backslash pi(s\_t,\; a\_t)\; -\; V^\backslash pi(s\_t))\backslash right]$$

which can be written as: $$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot A^{{\pi }_{\theta }}(s_t,a_t)]}}& & \text{(V.IV)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; A^\{\backslash pi\_\backslash theta\}(s\_t,\; a\_t)\backslash right]\}\; \backslash tag\{V.IV\}$$

and for sample-based approximation:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathrm{\nabla }}_{\theta }J(\theta )\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_{i,t}\mathrm{\mid }{s}_{i,t})\cdot {\widehat{A}}_{i,t}}}& & \text{(V.V)}\end{array}\backslash boxed\{\backslash nabla\_\backslash theta\; J(\backslash theta)\; \backslash approx\; \backslash frac\{1\}\{N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_\{i,t\}\; |\; s\_\{i,t\})\; \backslash cdot\; \backslash hat\{A\}\_\{i,t\}\}\; \backslash tag\{V.V\}$$

where ${\widehat{A}}_{i,t}\backslash hat\{A\}\_\{i,t\}$ is an estimate of the advantage function at time $tt$ in trajectory $ii$. This is the form of the policy gradient often used.

In practice, $A^{\pi }(s_t,a_t)A^\backslash pi(s\_t,\; a\_t)$ can be estimated as follows:

1. Learn a value function: Train a neural network $V_{\varphi }(s)V\_\backslash phi(s)$ (often called the "critic" or "value head") to approximate $V^{\pi }(s)V^\backslash pi(s)$. In LLM fine-tuning, this is often a linear layer on top of the same transformer backbone used for the policy.

2. Estimate $Q^{\pi }Q^\backslash pi$ from samples: Given a trajectory, the rewards-to-go ${\widehat{R}}_t={\sum }_{t^{\mathrm{\prime }}=t}^T{\gamma }^{t^{\mathrm{\prime }}}{r}_{t^{\mathrm{\prime }}}\backslash hat\{R\}\_t\; =\; \backslash sum\_\{t\text{'}=t\}^\{T\}\; \backslash gamma^\{t\text{'}\}\; r\_\{t\text{'}\}$ provides an unbiased (but high-variance) estimate of $Q^{\pi }(s_t,a_t)Q^\backslash pi(s\_t,\; a\_t)$.

3. Compute advantage estimates: ${\widehat{A}}_t={\widehat{R}}_t-V_{\varphi }(s_t)\backslash hat\{A\}\_t\; =\; \backslash hat\{R\}\_t\; -\; V\_\backslash phi(s\_t)$

More sophisticated methods like Generalized Advantage Estimation (GAE) interpolate between high-variance, low-bias estimates and low-variance, high-bias estimates by using a weighted combination of multi-step returns. See the GAE paper for more details.

VI: Importance Sampling and Off-Policy Policy Gradients

Note: In RL literature, "off-policy" typically refers to methods where the behavior policy (generating data) is arbitrarily quite different from the target policy (being optimized) say where transitions from policies thousands of updates old are reused. In this section, what we will call "off-policy" should more precisely be called "local off-policy".

The advantage-weighted policy gradient (V.IV) requires trajectories sampled from the current policy ${\pi }_{\theta }\backslash pi\_\backslash theta$. ... The advantage-weighted policy gradient (V.IV) requires trajectories sampled from the current policy ${\pi }_{\theta }\backslash pi\_\backslash theta$. This creates a fundamental inefficiency i.e., after each gradient update $\theta \to {\theta }^{\mathrm{\prime }}\backslash theta\; \backslash to\; \backslash theta\text{'}$ all previously collected trajectories become "stale" and we must discard these trajectories and sample new ones from the updated policy.

For LLMs, where each trajectory requires a full forward pass through billion(s)-parameter model, this is prohibitively expensive especially when we need many small gradient steps to train effectively.

We need a way to reuse the same trajectories for multiple gradient updates. Importance sampling provides the mathematical machinery to do exactly this!

Importance Sampling

Importance sampling is a technique for estimating expectations under one probability distribution using samples drawn from a different distribution. Consider an expectation for distribution $p(x)p(x)$:

$${\mathbb{E}}_{x\sim p}[f(x)]=\int p(x)f(x)dx\backslash mathbb\{E\}\_\{x\; \backslash sim\; p\}[f(x)]\; =\; \backslash int\; p(x)\; f(x)\; \backslash ,\; dx$$

We can rewrite this by multiplying and dividing by another distribution $q(x)q(x)$ (with $q(x)>0q(x)\; >\; 0$ wherever $p(x)>0p(x)\; >\; 0$ ):

$$=\int q(x)\frac{p(x)}{q(x)}f(x)dx={\mathbb{E}}_{x\sim q}[\frac{p(x)}{q(x)}f(x)]=\; \backslash int\; q(x)\; \backslash frac\{p(x)\}\{q(x)\}\; f(x)\; \backslash ,\; dx\; =\; \backslash mathbb\{E\}\_\{x\; \backslash sim\; q\}\backslash left[\backslash frac\{p(x)\}\{q(x)\}\; f(x)\backslash right]$$

The ratio $\frac{p(x)}{q(x)}\backslash frac\{p(x)\}\{q(x)\}$ is called the importance weight. This identity tells us:

$$\begin{array}{cccc}& \boxed{{\displaystyle {\mathbb{E}}_{x\sim p}[f(x)]={\mathbb{E}}_{x\sim q}[\frac{p(x)}{q(x)}f(x)]}}& & \text{(VI.I)}\end{array}\backslash boxed\{\backslash mathbb\{E\}\_\{x\; \backslash sim\; p\}[f(x)]\; =\; \backslash mathbb\{E\}\_\{x\; \backslash sim\; q\}\backslash left[\backslash frac\{p(x)\}\{q(x)\}\; f(x)\backslash right]\}\; \backslash tag\{VI.I\}$$

We can now estimate the expectation under $pp$ using samples from $qq$ as long as we reweight each sample by the ratio of probabilities.

Applying Importance Sampling to Policy Gradients

We can apply this technique to the policy gradient setting. The on-policy advantage-weighted gradient (V.IV) is:

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{\tau \sim {\pi }_{\theta }}[\sum _{t=0}^T{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot A^{{\pi }_{\theta }}(s_t,a_t)]\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{\backslash tau\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\backslash sum\_\{t=0\}^\{T\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; A^\{\backslash pi\_\backslash theta\}(s\_t,\; a\_t)\backslash right]$$

To apply importance sampling, we work at time-step level rather than trajectory level (full trajectory importance weights have extremely high variance). For a single timestep: $${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{(s_t,a_t)\sim {\pi }_{\theta }}[{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot A^{{\pi }_{\theta }}(s_t,a_t)]\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{(s\_t,\; a\_t)\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; A^\{\backslash pi\_\backslash theta\}(s\_t,\; a\_t)\backslash right]$$

Using importance sampling with samples from ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$:

$$={\mathbb{E}}_{(s_t,a_t)\sim {\pi }_{{\theta }_{\text{old}}}}[\frac{{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mathrm{\mid }{s}_t)}{\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\cdot A^{{\pi }_{\theta }}(s_t,a_t)]=\; \backslash mathbb\{E\}\_\{(s\_t,\; a\_t)\; \backslash sim\; \backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}\}\backslash left[\backslash frac\{\backslash pi\_\backslash theta(a\_t\; |\; s\_t)\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_t\; |\; s\_t)\}\; \backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\; \backslash cdot\; A^\{\backslash pi\_\backslash theta\}(s\_t,\; a\_t)\backslash right]$$

Now we apply the log-derivative identity ${\mathrm{\nabla }}_{\theta }\log {\pi }_{\theta }=\frac{{\mathrm{\nabla }}_{\theta }{\pi }_{\theta }}{{\pi }_{\theta }}\backslash nabla\_\backslash theta\; \backslash log\; \backslash pi\_\backslash theta\; =\; \backslash frac\{\backslash nabla\_\backslash theta\; \backslash pi\_\backslash theta\}\{\backslash pi\_\backslash theta\}$, which gives us a surrogate objective $L(\theta )L(\backslash theta)$ whose gradient equals this importance-weighted policy gradient:

$${\mathrm{\nabla }}_{\theta }J(\theta )={\mathbb{E}}_{(s_t,a_t)\sim {\pi }_{{\theta }_{\text{old}}}}[\frac{{\mathrm{\nabla }}_{\theta }{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mathrm{\mid }{s}_t)}{A}^{{\pi }_{{\theta }_{\text{old}}}}(s_t,a_t)]\backslash nabla\_\backslash theta\; J(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{(s\_t,\; a\_t)\; \backslash sim\; \backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}\}\backslash left[\backslash frac\{\backslash nabla\_\backslash theta\; \backslash pi\_\backslash theta(a\_t\; |\; s\_t)\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_t\; |\; s\_t)\}\; A^\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}\}(s\_t,\; a\_t)\backslash right]$$

where the importance-weighted surrogate objective also known as the Conservative Policy Iteration (CPI) objective is: $$L^{\text{CPI}}(\theta )={\mathbb{E}}_{(s_t,a_t)\sim {\pi }_{{\theta }_{\text{old}}}}[\frac{{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mathrm{\mid }{s}_t)}{A}^{{\pi }_{{\theta }_{\text{old}}}}(s_t,a_t)]L^\{\backslash text\{CPI\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_\{(s\_t,\; a\_t)\; \backslash sim\; \backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}\}\backslash left[\backslash frac\{\backslash pi\_\backslash theta(a\_t\; |\; s\_t)\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_t\; |\; s\_t)\}\; A^\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}\}(s\_t,\; a\_t)\backslash right]$$

We also define the probability ratio as:

$$\begin{array}{cccc}& r_t(\theta )=\frac{{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mathrm{\mid }{s}_t)}& & \text{(VI.II)}\end{array}r\_t(\backslash theta)\; =\; \backslash frac\{\backslash pi\_\backslash theta(a\_t\; |\; s\_t)\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_t\; |\; s\_t)\}\; \backslash tag\{VI.II\}$$

Note that $r_t({\theta }_{\text{old}})=1r\_t(\backslash theta\_\{\backslash text\{old\}\})\; =\; 1$ by construction. Thus, the CPI objective can be written as:

$$\begin{array}{cccc}& \boxed{{\displaystyle L^{\text{CPI}}(\theta )={\mathbb{E}}_t\left[\frac{{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mathrm{\mid }{s}_t)}{\widehat{A}}_t\right]={\mathbb{E}}_t[r_t(\theta ){\widehat{A}}_t]}}& & \text{(VI.III)}\end{array}\backslash boxed\{L^\{\backslash text\{CPI\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_t\backslash left[\backslash frac\{\backslash pi\_\backslash theta(a\_t\; |\; s\_t)\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_t\; |\; s\_t)\}\; \backslash hat\{A\}\_t\backslash right]\; =\; \backslash mathbb\{E\}\_t\backslash left[r\_t(\backslash theta)\; \backslash hat\{A\}\_t\backslash right]\}\; \backslash tag\{VI.III\}$$

where ${\widehat{A}}_t\backslash hat\{A\}\_t$ is the estimated advantage at timestep $tt$, and ${\mathbb{E}}_t[\cdot ]\backslash mathbb\{E\}\_t[\backslash cdot]$ denotes the empirical average over a batch of samples collected under ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$.

This objective has a clear interpretation:

• If ${\widehat{A}}_t>0\backslash hat\{A\}\_t\; >\; 0$ (action better than average), we want to increase $r_t(\theta )r\_t(\backslash theta)$, i.e., make the new policy more likely to take this action.
• If (action worse than average), we want to decrease $r_t(\theta )r\_t(\backslash theta)$, i.e., make the new policy less likely to take this action.

The corresponding sample-based approximation is:

$$\begin{array}{cccc}& \boxed{{\displaystyle L^{\text{CPI}}(\theta )\approx \frac{1}{N}\sum _{i=1}^N\sum _{t=0}^T\frac{{\pi }_{\theta }(a_{i,t}\mathrm{\mid }{s}_{i,t})}{{\pi }_{{\theta }_{\text{old}}}(a_{i,t}\mathrm{\mid }{s}_{i,t})}{\widehat{A}}_{i,t}}}& & \text{(VI.IV)}\end{array}\backslash boxed\{L^\{\backslash text\{CPI\}\}(\backslash theta)\; \backslash approx\; \backslash frac\{1\}\{N\}\; \backslash sum\_\{i=1\}^\{N\}\; \backslash sum\_\{t=0\}^\{T\}\; \backslash frac\{\backslash pi\_\backslash theta(a\_\{i,t\}\; |\; s\_\{i,t\})\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_\{i,t\}\; |\; s\_\{i,t\})\}\; \backslash hat\{A\}\_\{i,t\}\}\; \backslash tag\{VI.IV\}$$

Off-Policy Learning: Reusing Trajectories

The CPI objective enables off-policy learning: we can sample trajectories from ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$, store them and then perform multiple gradient updates on $\theta \backslash theta$ using the same batch of data. The typical workflow becomes:

1. Collect: Sample trajectories $\{{\tau }_i\}\backslash \{\backslash tau\_i\backslash \}$ from the current policy ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$
2. Compute: Calculate advantages ${\widehat{A}}_{i,t}\backslash hat\{A\}\_\{i,t\}$ and log-probabilities $\log {\pi }_{{\theta }_{\text{old}}}(a_{i,t}\mathrm{\mid }{s}_{i,t})\backslash log\; \backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_\{i,t\}\; |\; s\_\{i,t\})$
3. Store: Save the trajectories along with their advantages and old log-probabilities
4. Optimize: Perform multiple gradient ascent steps on $L^{\text{CPI}}(\theta )L^\{\backslash text\{CPI\}\}(\backslash theta)$ using mini-batches from the stored data
5. Repeat: Set ${\theta }_{\text{old}}\leftarrow \theta \backslash theta\_\{\backslash text\{old\}\}\; \backslash leftarrow\; \backslash theta$ and return to step 1

This dramatically improves sample efficiency. Instead of discarding trajectories after a single gradient step, we can extract multiple updates from each batch of expensive LLM rollouts.

The Instability Problem

While the CPI objective improves sample efficiency, unconstrained optimization of $L^{\text{CPI}}(\theta )L^\{\backslash text\{CPI\}\}(\backslash theta)$ is unstable. The core issue is that importance sampling becomes unreliable when ${\pi }_{\theta }\backslash pi\_\backslash theta$ drifts far from ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$:

• Extreme probability ratios: The ratio $r_t(\theta )r\_t(\backslash theta)$ can become arbitrarily large or small, destabilizing gradient estimates.
• Stale advantages: The estimates ${\widehat{A}}_t\backslash hat\{A\}\_t$ were computed under ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$ and become inaccurate as ${\pi }_{\theta }\backslash pi\_\backslash theta$ diverges. The optimizer may exploit these stale estimates, making updates that appear beneficial but are actually harmful.

In practice, unconstrained maximization of $L^{\text{CPI}}(\theta )L^\{\backslash text\{CPI\}\}(\backslash theta)$ often leads to excessively large policy updates that cause catastrophic performance collapse.

LLM Context (from Umar Jamil): Suppose we have a trajectory where the model generated "Shanghai is in China" with high advantage. Unconstrained optimization might dramatically upweight "China" as the next token given "Shanghai is in"—but this could simultaneously cause unintended probability shifts elsewhere, perhaps making the model overly likely to say "China" in completely unrelated contexts, or disrupting the probability mass across the entire vocabulary in unpredictable ways.

We need a mechanism to constrain ${\pi }_{\theta }\backslash pi\_\backslash theta$ from deviating too far from ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$ and keeping the ratio $r_t(\theta )r\_t(\backslash theta)$ close to 1 while still allowing meaningful policy improvement.

VII: Trust Region Policy Optimization (TRPO)

The CPI objective is attractive because it lets us reuse data via importance ratios, but unconstrained optimization is unstable. When ${\pi }_{\theta }\backslash pi\_\backslash theta$ drifts far from ${\pi }_{{\theta }_{\text{old}}}\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}$, the probability ratios $r_t(\theta )r\_t(\backslash theta)$ become extreme and the advantage estimates ${\widehat{A}}_t\backslash hat\{A\}\_t$ become stale and can be exploited by the optimizer.

The key insight of Trust Region Policy Optimization (TRPO) is that the surrogate objective $L^{\text{CPI}}(\theta )L^\{\backslash text\{CPI\}\}(\backslash theta)$ is only a valid approximation to the true objective within a local neighborhood of ${\theta }_{\text{old}}\backslash theta\_\{\backslash text\{old\}\}$. TRPO paper formalized this by proving policy performance is guaranteed to improve as long as the KL divergence between consecutive policies remains bounded. This theoretical result motivates constraining the policy update to stay within a "trust region" where the surrogate objective remains reliable. See the TRPO paper for the formal proof.

TRPO converts this insight into a constrained optimization problem that ensures the policy update stays within a "trust region" where the surrogate objective remains reliable.

The hyperparameter $\delta \backslash delta$ defines the trust region size, the maximum allowed divergence between consecutive policies. This constraint ensures that $r_t(\theta )r\_t(\backslash theta)$ remains close to 1, keeping our importance-weighted estimates reliable.

Solving (VII.I) requires second-order optimization. TRPO approximates the objective linearly and the KL constraint quadratically (using the Fisher Information Matrix) and then solves the resulting problem via the conjugate gradient algorithm followed by a line search to ensure constraints are satisfied.

For large-scale LLM training, this approach is impractical:

• Computational overhead: Each policy update requires multiple conjugate gradient iterations and line search steps, significantly more expensive than standard gradient descent.
• Memory requirements: Computing Fisher-vector products adds substantial memory overhead for billion(s)-parameter models

The theory behind TRPO also suggests using a KL penalty rather than a hard constraint. It is easier to implement and more computationally efficient.

$$\begin{array}{cccc}& \underset{\theta }{\max }{\mathbb{E}}_t[r_t(\theta ){\widehat{A}}_t-\beta \cdot D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}(\cdot \mathrm{\mid }{s}_t)\mathrm{\parallel }{\pi }_{\theta }(\cdot \mathrm{\mid }{s}_t))]& & \text{(VII.II)}\end{array}\backslash max\_\backslash theta\; \backslash ;\; \backslash mathbb\{E\}\_t\backslash left[r\_t(\backslash theta)\; \backslash hat\{A\}\_t\; -\; \backslash beta\; \backslash cdot\; D\_\{\backslash text\{KL\}\}\backslash left(\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(\backslash cdot|s\_t)\; \backslash |\; \backslash pi\_\backslash theta(\backslash cdot|s\_t)\backslash right)\backslash right]\; \backslash tag\{VII.II\}$$

However, choosing a penalty coefficient $\beta \backslash beta$ that works across different problems or even across different training stages is notoriously difficult. This motivates Proximal Policy Optimization (PPO): a first-order method that achieves TRPO's stability through a clipped surrogate objective rather than explicit constraints.

VIII: Proximal Policy Optimization (PPO)

Proximal Policy Optimization (PPO) achieves TRPO's stability guarantees using only first-order optimization. Instead of explicitly constraining the KL divergence, PPO modifies the objective function itself to discourage large policy updates through a clipping mechanism. It implicitly limits how far the policy can move, providing a "soft" trust region using only standard gradient descent.

Clipped Surrogate Objective

CPI objective and probability ratio from Section VI:

$$L^{\text{CPI}}(\theta )={\mathbb{E}}_t[r_t(\theta ){\widehat{A}}_t]\text{where}{r}_t(\theta )=\frac{{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mathrm{\mid }{s}_t)}L^\{\backslash text\{CPI\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_t\backslash left[r\_t(\backslash theta)\; \backslash hat\{A\}\_t\backslash right]\; \backslash quad\; \backslash text\{where\}\; \backslash quad\; r\_t(\backslash theta)\; =\; \backslash frac\{\backslash pi\_\backslash theta(a\_t\; |\; s\_t)\}\{\backslash pi\_\{\backslash theta\_\{\backslash text\{old\}\}\}(a\_t\; |\; s\_t)\}$$

The problem with $L^{\text{CPI}}L^\{\backslash text\{CPI\}\}$ is that nothing prevents $r_t(\theta )r\_t(\backslash theta)$ from becoming arbitrarily large or small. PPO addresses this by clipping the probability ratio to stay within $[1-ϵ,1+ϵ][1-\backslash epsilon,\; 1+\backslash epsilon]$:

$$\begin{array}{cccc}& \boxed{{\displaystyle L^{\text{CLIP}}(\theta )={\mathbb{E}}_t[\min (r_t(\theta ){\widehat{A}}_t,\text{clip}(r_t(\theta ),1-ϵ,1+ϵ)\cdot {\widehat{A}}_t)]}}& & \text{(VIII.I)}\end{array}\backslash boxed\{L^\{\backslash text\{CLIP\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_t\backslash left[\backslash min\backslash left(r\_t(\backslash theta)\; \backslash hat\{A\}\_t,\; \backslash ;\; \backslash text\{clip\}(r\_t(\backslash theta),\; 1-\backslash epsilon,\; 1+\backslash epsilon)\; \backslash cdot\; \backslash hat\{A\}\_t\backslash right)\backslash right]\}\; \backslash tag\{VIII.I\}$$

where $ϵ\backslash epsilon$ is a hyperparameter ( $ϵ=0.2\backslash epsilon\; =\; 0.2$ from the PPO paper) and the clip function is defined as:

The $\min \backslash min$ operator in (VIII.I) is important. It ensures we take the more pessimistic (lower) estimate between the clipped and unclipped objectives. This creates different behavior depending on the sign of the advantage:

Case 1: Positive Advantage ( ${\widehat{A}}_t>0\backslash hat\{A\}\_t\; >\; 0$ )

When an action is better than average, we want to increase its probability, which means increasing $r_t(\theta )r\_t(\backslash theta)$. The objective becomes:

$$L_t^{\text{CLIP}}=\min (r_t(\theta ),1+ϵ)\cdot {\widehat{A}}_{t}L^\{\backslash text\{CLIP\}\}\_t\; =\; \backslash min\backslash left(r\_t(\backslash theta),\; 1+\backslash epsilon\backslash right)\; \backslash cdot\; \backslash hat\{A\}\_t$$

• If $r_t(\theta )\le 1+ϵr\_t(\backslash theta)\; \backslash leq\; 1+\backslash epsilon$: The objective is $r_t(\theta ){\widehat{A}}_{t}r\_t(\backslash theta)\; \backslash hat\{A\}\_t$, so gradient ascent increases $r_t(\theta )r\_t(\backslash theta)$
• If $r_t(\theta )>1+ϵr\_t(\backslash theta)\; >\; 1+\backslash epsilon$: The objective becomes $(1+ϵ){\widehat{A}}_t(1+\backslash epsilon)\backslash hat\{A\}\_t$

The clipping removes the incentive to increase $r_t(\theta )r\_t(\backslash theta)$ beyond $1+ϵ1+\backslash epsilon$.

Case 2: Negative Advantage ( )

When an action is worse than average, we want to decrease its probability, which means decreasing $r_t(\theta )r\_t(\backslash theta)$. Since , multiplying by a smaller $r_{t}r\_t$ makes the product less negative (larger). The objective becomes:

$$L_t^{\text{CLIP}}=\max (r_t(\theta ),1-ϵ)\cdot {\widehat{A}}_{t}L^\{\backslash text\{CLIP\}\}\_t\; =\; \backslash max\backslash left(r\_t(\backslash theta),\; 1-\backslash epsilon\backslash right)\; \backslash cdot\; \backslash hat\{A\}\_t$$

(The $\min \backslash min$ with negative values becomes a $\max \backslash max$ in terms of which $r_{t}r\_t$ is selected.)

• If $r_t(\theta )\ge 1-ϵr\_t(\backslash theta)\; \backslash geq\; 1-\backslash epsilon$: The objective is $r_t(\theta ){\widehat{A}}_{t}r\_t(\backslash theta)\; \backslash hat\{A\}\_t$, so gradient ascent decreases $r_t(\theta )r\_t(\backslash theta)$
• If : The objective becomes $(1-ϵ){\widehat{A}}_t(1-\backslash epsilon)\backslash hat\{A\}\_t$

The clipping removes the incentive to decrease $r_t(\theta )r\_t(\backslash theta)$ beyond $1-ϵ1-\backslash epsilon$.

The takeaway here is that PPO provides a pessimistic lower bound on $L^{\text{CPI}}L^\{\backslash text\{CPI\}\}$. We ignore updates when they would make things "too good to be true."

LLM Context (from Umar Jamil Video): In language model fine-tuning, the policy ${\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)\backslash pi\_\backslash theta(a\_t|s\_t)$ is the probability the model assigns to token $a_{t}a\_t$ given the context $s_{t}s\_t$ (prompt + previously generated tokens). The probability ratio $r_t(\theta )r\_t(\backslash theta)$ measures how much more or less likely the fine-tuned model is to generate a particular token compared to the reference policy. Clipping ensures that no single token's probability can change by more than a factor of $(1\pm ϵ)(1\; \backslash pm\; \backslash epsilon)$ in a single update iteration, preventing the model from "overreacting" to high-advantage tokens.

PPO Objective

In practice, PPO combines the clipped policy objective with two additional terms:

$$\begin{array}{cccc}& \boxed{{\displaystyle L^{\text{PPO}}(\theta )={\mathbb{E}}_t[L_t^{\text{CLIP}}(\theta )-c_1{L}_t^{\text{VF}}(\theta )+c_{2}S[{\pi }_{\theta }](s_t)]}}& & \text{(VIII.II)}\end{array}\backslash boxed\{L^\{\backslash text\{PPO\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_t\backslash left[L^\{\backslash text\{CLIP\}\}\_t(\backslash theta)\; -\; c\_1\; L^\{\backslash text\{VF\}\}\_t(\backslash theta)\; +\; c\_2\; S[\backslash pi\_\backslash theta](s\_t)\backslash right]\}\; \backslash tag\{VIII.II\}$$

1. Value Function Loss ( $L^{\text{VF}}L^\{\backslash text\{VF\}\}$ ): Recall from Section V that we need a value function $V_{\varphi }(s)V\_\backslash phi(s)$ to compute advantage estimates. The value function is trained to minimize the squared error between its predictions and the actual returns:

$$L_t^{\text{VF}}(\theta )={(V_{\theta }(s_t)-V_t^{\text{target}})}^{2}L^\{\backslash text\{VF\}\}\_t(\backslash theta)\; =\; \backslash left(V\_\backslash theta(s\_t)\; -\; V\_t^\{\backslash text\{target\}\}\backslash right)^2$$

where $V_t^{\text{target}}V\_t^\{\backslash text\{target\}\}$ is typically the discounted return-to-go. When the policy and value function share parameters (common in LLM fine-tuning where both use the same transformer backbone), this loss is subtracted from the objective (hence the negative sign, since we maximize $L^{\text{PPO}}L^\{\backslash text\{PPO\}\}$ but minimize $L^{\text{VF}}L^\{\backslash text\{VF\}\}$ ).

2. Entropy Bonus ( $S[{\pi }_{\theta }]S[\backslash pi\_\backslash theta]$ ): To encourage exploration and prevent premature convergence to deterministic policies, PPO adds an entropy loss:

$$S[{\pi }_{\theta }](s_t)=-\sum _a{\pi }_{\theta }(a\mathrm{\mid }{s}_t)\log {\pi }_{\theta }(a\mathrm{\mid }{s}_t)S[\backslash pi\_\backslash theta](s\_t)\; =\; -\backslash sum\_a\; \backslash pi\_\backslash theta(a|s\_t)\; \backslash log\; \backslash pi\_\backslash theta(a|s\_t)$$

Here, the coefficients $c_1,c_2>0c\_1,\; c\_2\; >\; 0$ control the regularization strength.

IX: Complete PPO Objective with KL Penalty

When fine-tuning an LLM with "vanilla" PPO, the policy learns to maximize rewards from the reward model. However, the reward model is an imperfect proxy for human preferences. It is a neural network trained on limited data that can be exploited. Without constraints, the policy may discover adversarial outputs that achieve high reward scores while producing text that:

• Degenerates into repetitive or nonsensical patterns that "fool" the reward model
• Drifts far from natural language, losing fluency and coherence
• Exploits spurious correlations learned by the reward model

This phenomenon is called reward hacking. The policy finds a way to "game" the reward model rather than genuinely improving response quality.

To prevent reward hacking, the InstructGPT paper adds a KL divergence penalty that regularizes the policy to stay close to a reference model ${\pi }_{\text{ref}}\backslash pi\_\{\backslash text\{ref\}\}$ (typically the SFT model before RL fine-tuning).

From Section VIII, the PPO objective (to be maximized via gradient ascent) consists of three terms:

$$L^{\text{PPO}}(\theta )=\underset{\text{Clipped Policy Objective}}{\underbrace{L^{\text{CLIP}}(\theta )}}-\underset{\text{Value Function Loss}}{\underbrace{c_1{L}^{\text{VF}}(\theta )}}+\underset{\text{Entropy Bonus}}{\underbrace{c_{2}S[{\pi }_{\theta }]}}L^\{\backslash text\{PPO\}\}(\backslash theta)\; =\; \backslash underbrace\{L^\{\backslash text\{CLIP\}\}(\backslash theta)\}\_\{\backslash text\{Clipped\; Policy\; Objective\}\}\; -\; \backslash underbrace\{c\_1\; L^\{\backslash text\{VF\}\}(\backslash theta)\}\_\{\backslash text\{Value\; Function\; Loss\}\}\; +\; \backslash underbrace\{c\_2\; S[\backslash pi\_\backslash theta]\}\_\{\backslash text\{Entropy\; Bonus\}\}$$

Now, we don't use raw reward model scores directly. Instead, we define a KL-penalized reward that regularizes the policy to stay close to a reference model ${\pi }_{\text{ref}}\backslash pi\_\{\backslash text\{ref\}\}$:

$$\begin{array}{cccc}& \boxed{{\displaystyle r_{\text{total}}(s_t,a_t)=r_{\text{RM}}(s_t,a_t)-\beta \cdot D_{\text{KL}}({\pi }_{\theta }(\cdot \mathrm{\mid }{s}_t)\mathrm{\parallel }{\pi }_{\text{ref}}(\cdot \mathrm{\mid }{s}_t))}}& & \text{(IX.I)}\end{array}\backslash boxed\{r\_\{\backslash text\{total\}\}(s\_t,\; a\_t)\; =\; r\_\{\backslash text\{RM\}\}(s\_t,\; a\_t)\; -\; \backslash beta\; \backslash cdot\; D\_\{\backslash text\{KL\}\}\backslash left(\backslash pi\_\backslash theta(\backslash cdot|s\_t)\; \backslash |\; \backslash pi\_\{\backslash text\{ref\}\}(\backslash cdot|s\_t)\backslash right)\}\; \backslash tag\{IX.I\}$$

where:

• $r_{\text{RM}}(s_t,a_t)r\_\{\backslash text\{RM\}\}(s\_t,\; a\_t)$ is the reward signal at timestep $tt$
• $\beta \backslash beta$ is the KL penalty coefficient
• ${\pi }_{\text{ref}}\backslash pi\_\{\backslash text\{ref\}\}$ is the frozen reference model

At each token position, the KL divergence simplifies to:

$$D_{\text{KL}}({\pi }_{\theta }(\cdot \mathrm{\mid }{s}_t)\mathrm{\parallel }{\pi }_{\text{ref}}(\cdot \mathrm{\mid }{s}_t))={\mathbb{E}}_{a\sim {\pi }_{\theta }}[\log \frac{{\pi }_{\theta }(a\mathrm{\mid }{s}_t)}{{\pi }_{\text{ref}}(a\mathrm{\mid }{s}_t)}]D\_\{\backslash text\{KL\}\}\backslash left(\backslash pi\_\backslash theta(\backslash cdot|s\_t)\; \backslash |\; \backslash pi\_\{\backslash text\{ref\}\}(\backslash cdot|s\_t)\backslash right)\; =\; \backslash mathbb\{E\}\_\{a\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[\backslash log\; \backslash frac\{\backslash pi\_\backslash theta(a|s\_t)\}\{\backslash pi\_\{\backslash text\{ref\}\}(a|s\_t)\}\backslash right]$$

In practice we estimate this expectation with the sampled token $a_{t}a\_t$, yielding: $${\widehat{d}}_t=\log \frac{{\pi }_{\theta }(a_t\mathrm{\mid }{s}_t)}{{\pi }_{\mathrm{r}\mathrm{e}\mathrm{f}}(a_t\mathrm{\mid }{s}_t)}\backslash hat\; d\_t=\backslash log\; \backslash frac\{\backslash pi\_\backslash theta(a\_t|s\_t)\}\{\backslash pi\_\{\backslash mathrm\{ref\}\}(a\_t|s\_t)\}$$

Note that the reward model $r_{\varphi }(x,y)r\_\backslash phi(x,\; y)$ produces a single scalar for the complete response $(x,y)(x,\; y)$. This score is assigned only at the final token $TT$, while the KL penalty applies at every token.

The KL penalty serves two purposes:

1. Prevents reward hacking: The policy cannot drift arbitrarily far from natural language
2. Maintains fluency: Outputs remain similar in distribution to the well-trained SFT model

It modifies the advantage estimates ${\widehat{A}}_t\backslash hat\{A\}\_t$ used in PPO through the modified per-token rewards. However, it is mathematically equivalent (and more efficient in implementation) to add the KL term directly to the objective. The PPO objective with KL penalty is:

$$J(\theta )=\underset{\text{Vanilla PPO objective}}{\underbrace{{\mathbb{E}}_{a\sim {\pi }_{\theta }}[r_{\text{RM}}(s,a)]}}-\underset{\text{KL penalty term}}{\underbrace{\beta \cdot D_{\text{KL}}({\pi }_{\theta }\mathrm{\parallel }{\pi }_{\text{ref}})}}J(\backslash theta)\; =\; \backslash underbrace\{\backslash mathbb\{E\}\_\{a\; \backslash sim\; \backslash pi\_\backslash theta\}\backslash left[r\_\{\backslash text\{RM\}\}(s,\; a)\backslash right]\}\_\{\backslash text\{Vanilla\; PPO\; objective\}\}\; -\; \backslash underbrace\{\backslash beta\; \backslash cdot\; D\_\{\backslash text\{KL\}\}(\backslash pi\_\backslash theta\; \backslash |\; \backslash pi\_\{\backslash text\{ref\}\})\}\_\{\backslash text\{KL\; penalty\; term\}\}$$

The first term is exactly what vanilla PPO optimizes using the clipped surrogate. The KL penalty term appears as a separate additive component that penalizes divergence from the reference model. Substituting the PPO clipped surrogate for the first term:

$$J_{\text{c}}(\theta )={\mathbb{E}}_t[\min (r_t(\theta ){\widehat{A}}_t,\text{clip}(r_t(\theta ),1-ϵ,1+ϵ){\widehat{A}}_t)]-\beta \cdot D_{\text{KL}}({\pi }_{\theta }\mathrm{\parallel }{\pi }_{\text{ref}})J\_\{\backslash text\{c\}\}(\backslash theta)\; =\; \backslash mathbb\{E\}\_t\backslash left[\backslash min\backslash left(r\_t(\backslash theta)\; \backslash hat\{A\}\_t,\; \backslash text\{clip\}(r\_t(\backslash theta),\; 1-\backslash epsilon,\; 1+\backslash epsilon)\; \backslash hat\{A\}\_t\backslash right)\backslash right]\; -\; \backslash beta\; \backslash cdot\; D\_\{\backslash text\{KL\}\}(\backslash pi\_\backslash theta\; \backslash |\; \backslash pi\_\{\backslash text\{ref\}\})$$

Combining all components, the complete PPO objective with KL penalty (to be maximized) is:

$$\begin{array}{cccc}& \boxed{{\displaystyle L^{\text{RLHF}}(\theta )=\underset{\text{Policy Objective}}{\underbrace{L^{\text{CLIP}}(\theta )}}-\underset{\text{Value Loss}}{\underbrace{c_1{L}^{\text{VF}}(\theta )}}+\underset{\text{Entropy Bonus}}{\underbrace{c_{2}S[{\pi }_{\theta }]}}-\underset{\text{KL Penalty}}{\underbrace{\beta \cdot D_{\text{KL}}({\pi }_{\theta }\mathrm{\parallel }{\pi }_{\text{ref}})}}}}& & \text{(IX.II)}\end{array}\backslash boxed\{L^\{\backslash text\{RLHF\}\}(\backslash theta)\; =\; \backslash underbrace\{L^\{\backslash text\{CLIP\}\}(\backslash theta)\}\_\{\backslash text\{Policy\; Objective\}\}\; -\; \backslash underbrace\{c\_1\; L^\{\backslash text\{VF\}\}(\backslash theta)\}\_\{\backslash text\{Value\; Loss\}\}\; +\; \backslash underbrace\{c\_2\; S[\backslash pi\_\backslash theta]\}\_\{\backslash text\{Entropy\; Bonus\}\}\; -\; \backslash underbrace\{\backslash beta\; \backslash cdot\; D\_\{\backslash text\{KL\}\}(\backslash pi\_\backslash theta\; \backslash |\; \backslash pi\_\{\backslash text\{ref\}\})\}\_\{\backslash text\{KL\; Penalty\}\}\}\; \backslash tag\{IX.II\}$$

Here, each term serves a distinct purpose:

It is important to distinguish the two KL-related mechanisms in the complete loss. The PPO clipping mechanism acts as a short-term anchor that constrains how much the policy can change in a single update, while the KL penalty is a long-term anchor that constrains how far the policy can drift from its starting point across all of training.

Finally done...

And that's the full derivation! What I find satisfying is that every term in the final loss has a specific purpose. Each one exists because we ran into a specific problem along the way and needed to fix it. I will admit it was not easy to understand all the math and concepts behind the loss. I still do not fully understand every detail but I understand it far better than I did a few days ago.

I hope this was useful. If you spot any errors in derivation (which I'm sure there are) or have suggestions, feel free to reach out.

References

• Video:

  • Umar Jamil's video on RLHF and PPO: A comprehensive and must-watch video covering RLHF and PPO concepts.

• Papers:

  • Proximal Policy Optimization Algorithms: The foundational PPO paper introducing the clipped surrogate objective.
  • Training language models to follow instructions with human feedback: The InstructGPT paper demonstrating PPO with KL penalty to mitigate reward hacking in LLM fine-tuning.
  • Trust Region Policy Optimization: The TRPO paper that motivates the trust region constraints used in PPO.
  • High-Dimensional Continuous Control Using Generalized Advantage Estimation: GAE paper introducing the exponentially-weighted advantage estimator for variance reduction in policy gradients.