SERIES: Learning RL and JAX in Public - from zero to DeepMind :)
When people ask "how was ChatGPT trained?", the answer usually involves RLHF - Reinforcement Learning from Human Feedback. And the RL part of RLHF is PPO.
Day 8 is PPO. And it turns out to be simpler than I expected.
The problem that PPO solves
Actor-Critic (Day 7) worked, but it had an instability problem. A single episode with an unusually high or low return would push the network weights too far in one direction. The policy would change dramatically. Sometimes it would unlearn things it already knew. Training would collapse.
The fix is elegant: put a speed limit on learning.
The one idea that is PPO
After each update, PPO compares the new policy to the old one:
ratio = new_probability / old_probability
If ratio = 1.0: nothing changed for this action.
If ratio = 1.5: the new policy is 50% more likely to take this action.
If ratio = 0.5: the new policy is 50% less likely.
Then PPO clips this ratio. It says: I will not let you change by more than 20% in one update.
clipped_ratio = clip(ratio, 0.8, 1.2)
loss = -min(ratio * advantage, clipped_ratio * advantage)
That is literally it. One clip. One min. That is the entirety of what makes PPO different from Actor-Critic.
The gradient from a clipped update becomes zero once the ratio hits the boundary. The policy stops updating further for that action in that step. Come back next batch and nudge it again if you want more.
Small stable steps, every update.
The three things PPO adds over Actor-Critic
1. Clipping (the main idea)
Keeps the policy from changing too much in one shot.
2. Multiple epochs per batch
Actor-Critic uses each episode once. PPO collects a batch of episodes and then runs 4 gradient updates on that same data. More learning per episode. More sample efficient.
3. Entropy bonus
Entropy measures how spread out the action probabilities are. High entropy means the agent is still considering many options. Low entropy means it has collapsed to always picking one action.
PPO adds a small reward for entropy:
total_loss = policy_loss + critic_loss - 0.01 * entropy
This keeps the agent exploring longer before committing. Without it, policies can collapse to one action too early and get stuck.
What the training loop looks like
for iteration in range(200):
batch = collect_batch(actor, critic, episodes=10)
for epoch in range(4):
actor_loss = clipped_ppo_loss(batch)
critic_loss = value_loss(batch)
update(actor, critic)
Two loops. Outer loop collects data. Inner loop extracts learning from that data. Actor-Critic had one loop. This is the reason PPO trains faster per episode collected.
What I saw when I ran this
The value estimates from the critic after training:
states near goal: +0.6 to +0.8
safe middle path: +0.1 to +0.4
states near holes: -0.3 to -0.6
The critic mapped out the whole gridworld just by watching the actor fail and succeed. Nobody told it where the holes were. It figured it out.
The policy arrows lined up cleanly. Same result as Q-learning (Day 4), DQN (Day 5), REINFORCE (Day 6), Actor-Critic (Day 7). Different algorithm every time. Same learned behavior.
That pattern keeps hitting me. The Bellman intuition runs through all of it.
Why PPO became the default
Three reasons: it is stable (the clip prevents collapse), it is simple (one extra line over Actor-Critic), and it is general (discrete actions, continuous actions, LLM fine-tuning - all the same algorithm).
When OpenAI trained ChatGPT, they had humans rank responses. Those rankings became a reward signal. PPO optimized the language model against that signal. The RL loop you just implemented is the same loop, just with a very different environment and a very different reward function.
When I built the forge project (a GRPO trainer), I did not fully understand why it worked the way it did. GRPO is PPO with one change: instead of a single critic estimating advantage, it runs a group of episodes and compares their returns to each other. Same clipping. Same stability. The group comparison replaces the critic.
Five days ago I could not have explained that. Now I can.
Day 9: Haiku. DeepMind's neural network library. We rewrite all of this in a fraction of the code, and it starts looking like actual research-grade JAX.
All code from this series, organised by day, is on my GitHub: https://github.com/MadhumithaKolkar/jax-rl-lab
Happy learning everyone !
~ Madhumitha Kolkar (index_0)