JOVANA
Explore Library Glossary Getting Started Three Levels Fields How it works Mission
Join the mission
All guides

Optimizing the policy with PPO and a KL leash

Now you have a reward model. Run RL against it — but on a leash, so the policy improves the answer without wandering off into nonsense the reward model happens to love.

The RL step, finally

With a reward model standing in for human judgement, the third move is plain reinforcement learning: improve the policy to earn more reward. The workhorse here is PPO for preference optimization — using Proximal Policy Optimization (PPO) where the reward comes not from the environment but from the reward model you just trained. This combination is the classic PPO-for-RLHF loop.

The RL step at last — the policy acts and the reward model returns the reward that drives each update.

Diagram of the agent-environment reinforcement-learning loop: the agent takes an action and the environment returns a reward and a new state.

Why PPO and not something simpler? Because the reward is fragile. PPO improves the policy in small, conservative steps via its clipped surrogate objective, which refuses to move too far on any single update. When your reward is an imperfect learned model, small steps are a feature, not a limitation.

Why a leash is non-negotiable

Here is the crucial twist that separates RLHF from textbook RL. If you simply maximize the reward model, the policy will eventually discover inputs the reward model scores absurdly high but humans hate — gibberish that happens to trip its blind spots. The reward model is a proxy, and pushed hard enough every proxy breaks.

The fix is a KL penalty to a reference policy. We keep a frozen copy of the model from before RL — the reference — and penalize the policy whenever it drifts too far from it, measured by KL divergence. The optimizer is told: earn more reward, but don't become a different model to do it. This is the same idea as the broader KL penalty in alignment.

R(x,y)=r_\phi(x,y)-\beta\,\mathrm{KL}\!\left(\pi_\theta(\cdot\mid x)\,\|\,\pi_{\text{ref}}(\cdot\mid x)\right)

The leash as an equation: the effective reward is the model's score minus beta times the KL drift from the frozen reference policy.

The full loop in one place

Putting the three moves together, here is the complete RLHF training step that runs over and over:

  1. Sample a batch of prompts and let the current policy generate responses.
  2. Score each response with the reward model to get a reward.
  3. Subtract a KL term: reward minus β times the divergence from the reference policy.
  4. Update the policy with PPO's clipped objective on this penalized reward.
  5. Repeat — periodically refreshing comparisons so the reward model stays ahead of the policy.
L^{\mathrm{CLIP}}(\theta)=\mathbb{E}_t\!\left[\min\!\left(\rho_t(\theta)\,\hat{A}_t,\;\mathrm{clip}\!\left(\rho_t(\theta),\,1-\epsilon,\,1+\epsilon\right)\hat{A}_t\right)\right],\quad \rho_t(\theta)=\frac{\pi_\theta(a_t\mid s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t\mid s_t)}

PPO's clipped surrogate objective: the probability ratio is clipped to a small band so each policy update stays conservative.

# one RLHF reward, per response
r_rm   = reward_model(prompt, response)         # learned human-preference score
kl     = logp_policy(response) - logp_ref(response)  # drift from reference
reward = r_rm - beta * kl                        # reward on a leash
# ...then a normal PPO update using `reward`
The effective reward is the learned score minus a KL leash to the reference policy.

Knobs that actually matter

In practice three settings dominate. β (the KL weight) controls the leash and is the single most important dial. A small entropy bonus keeps the policy from collapsing to one phrasing too early, preserving exploration. And the number of PPO steps per batch trades stability for speed — more steps reuse data but risk overshooting the conservative region PPO is designed to respect.