RLHF Reward Model: KL-Penalized Alignment (2D)

Alignment Objective

KL penalty β (anchor strength) 0.30

Learning rate η 0.15

Live State

E[reward]

0.00

KL(π ‖ π_ref)

0.000

Entropy H(π)

0.00

Train step

Bar height = π(a) Dashed tick = reference π_ref(a)

How it works

RLHF fine-tunes a policy π over responses using a learned reward model r(a), but optimizing r(a) alone lets the policy drift arbitrarily far from the reference model π_ref and "reward-hack" — exploiting a blind spot the reward model scores highly but a human wouldn't. The standard fix (used in InstructGPT-style RLHF and DPO's implicit objective) adds a KL-divergence penalty against the reference policy:

maximize  J(π) = E_a~π[ r(a) ]  −  β · KL(π ‖ π_ref)

closed-form optimum:  π*(a) ∝ π_ref(a) · exp( r(a) / β )

This simulator runs gradient ascent on eight discrete "responses" arranged as 2D bars. Each step nudges the policy's logits toward higher reward while a term pulls them back toward the reference log-probabilities, scaled by β:

logit_a += η · [ r(a) − r̄ − β · log(π(a) / π_ref(a)) ]
π(a) = softmax(logits)

β (KL penalty) — low β lets the policy chase reward almost unconstrained (fast reward, big drift from π_ref); high β keeps it close to the reference at the cost of ignoring reward signal.
η (learning rate) — how fast the logits move each training step.
Simulate reward hack — spikes the reward of one low-quality response to mimic a reward-model exploit; watch whether the KL anchor (β) is strong enough to resist the policy collapsing onto it.
Bar height = π(a); the reference distribution π_ref is drawn as a dashed tick at its own height for comparison. Color ramps red→green with r(a).

Real-world relevance: this is the exact KL-regularized objective behind PPO-based RLHF and the reward term DPO optimizes implicitly without a separate reward model — β is the same "KL coefficient" hyperparameter tuned in production alignment pipelines.

⚠ Couldn't renderCanvas 2D context failed to initialize. Try reloading.