I Implemented the Algorithm Behind ChatGPT From Scratch - Day 8 (PPO).

A developer implementing reinforcement learning algorithms from scratch in JAX reports that PPO (Proximal Policy Optimization), the algorithm behind ChatGPT's RLHF training, is simpler than expected. The developer found that PPO's core innovation is a single clipping operation that limits policy updates, preventing training collapse, and that it adds multiple epochs per batch and an entropy bonus over standard Actor-Critic. The developer's experiments on a gridworld showed the critic successfully mapped the environment's safe and dangerous states, and the learned policy matched results from previous algorithms, highlighting the underlying Bellman intuition.

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 : collect 10 episodes with current policy batch = collect batch actor, critic, episodes=10 PPO epochs: squeeze 4 updates out of this batch 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 https://github.com/MadhumithaKolkar/jax-rl-lab Happy learning everyone ~ Madhumitha Kolkar index 0