RLHF · L3 Code — LLM Atlas

Implement the core pieces of a PPO update for language models: reward shaping and the clipped objective.

Policy model (trainable)

Reference model (frozen)

Reward model (frozen)

Value head (trainable)

Four models resident in GPU memory during PPO training.

You will rarely hand-write a full PPO trainer — libraries like TRL exist for this. But understanding the two functions below (reward shaping with KL, and the clipped surrogate loss) demystifies what those libraries do under the hood.