RLHF: Reward Modeling + KL-Regularized Policy Optimization

With π = softmax(z), log π(a) = z_a − log(e^z_A + e^z_B), so ∇_z log π(a) = e_a − π. For a record’s ratio, ∇_zρ_i = ρ_i(e_ai − π), because π_old(a_i) is a fixed number. At z_old every ρ_i = 1 lies strictly inside the band, so each term is ρ_iÂ_i nearby and ∇L = ¼ Σ Â_i(e_ai − π_old) = g. L is maximized, so the step adds ηg.

The full reward table is known here, so the true objective is available: J(π) = π(A)·1 + π(B)·0 = π(A), with ∂J/∂z_c = π(c)(R(c) − J). At π_old the true advantages are R(a) − J = (0.5, −0.5) and ∇J = (0.25, −0.25). The standardized weights of AABB, (+1, +1, −1, −1), are not these advantages. Because the table is known, sampling is unnecessary for solving this toy problem; the four records illustrate an estimator and a surrogate, not a need for reinforcement learning.

PPO usually estimates advantages across time steps with generalized advantage estimation (GAE), which bootstraps from a learned value function; the PPO reference page gives the formula. With one action, no later state and no value function, there is nothing to discount or bootstrap, and these weights come from the group, not from GAE.

Every order with the same count gives the same gradient, √(k(4 − k)) / 4 for A. Weighting the rows by their counts 1, 4, 6, 4, 1, the expected normalized gradient is 3/16 + √3/8 = 0.404006, not the true ∇J_A = 0.25. Averaging the five menu groups equally would give 0.273205, which weights them wrongly.

Raw rewards with no baseline average exactly to ∇J. Centering alone already gives 0.1875 = ¾ · 0.25: the group mean includes the record’s own reward, so this baseline is not independent of the action and the usual unbiasedness argument for a baseline does not apply. Dividing by a shared σ also makes the four weights depend on one another, although the draws are independent. These are exact finite calculations; a larger expected gradient is not evidence that this estimator learns better or worse.

Sources: Schulman et al. (2017). Proximal Policy Optimization Algorithms. arXiv:1707.06347v2. Shao et al. (2024). DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models. arXiv:2402.03300v1 (PDF). Hugging Face TRL v0.24.0: advantage normalization and loss. PyTorch documentation: torch.std.

A separate implementation of the validation, normalization, step, surrogate, candidates and sixteen orders, with NumPy float64 arrays and the same refusals. The site’s test suite runs this exact text and compares it with the TypeScript helper on every ordered group, every short binary score list, a grid of candidates and each refusal. Results agree to within a few units in the last place, not bit for bit.

import math
import re
from fractions import Fraction
from itertools import product

import numpy as np

# One context, one choice of A or B, then the episode ends: its return is that one reward.
ACTIONS = ("A", "B")
REWARD = {"A": 1, "B": 0}          # the rule chosen for this example, not a learned reward
Z_OLD = np.array([0.0, 0.0])       # frozen old logits, shape (2,), float64
ETA = 0.5                          # one ascent step: this objective is maximized
LOWER, UPPER = 0.8, 1.2            # the ratio band 1 - 0.2 to 1 + 0.2 (clip width 0.2)
ID_PATTERN = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,31}", re.ASCII)


class GroupUpdateError(ValueError):
    """A coded refusal. Nothing is repaired, zero-filled or dropped."""

    def __init__(self, code, detail):
        super().__init__(f"{code}: {detail}")
        self.code = code


def softmax(z):
    """Internal: two validated float64 logits of shape (2,), in shifted coordinates."""
    e = np.exp(z - z.max())
    return e / e.sum()


PI_OLD = softmax(Z_OLD)            # (0.5, 0.5): the denominator of every ratio, never re-estimated


def is_number(x):
    """A Python int or float (NumPy float64 is a float). A bool, a NumPy int or a string is not a score."""
    return isinstance(x, (int, float)) and not isinstance(x, bool)


def validate_group(records):
    """Four records, checked one at a time: shape, id, duplicate, action, then the reward."""
    if not isinstance(records, list) or len(records) != 4:
        raise GroupUpdateError("malformed-group", "a group is a list of exactly four records")
    seen, valid = set(), []
    for i, record in enumerate(records, start=1):
        if type(record) is not dict or set(record) != {"id", "action", "reward"}:
            raise GroupUpdateError("malformed-record", f"record {i} needs exactly id, action and reward")
        rid, action, reward = record["id"], record["action"], record["reward"]
        if type(rid) is not str or not ID_PATTERN.fullmatch(rid):
            raise GroupUpdateError("invalid-id", f"record {i} needs a string ID such as e1")
        if rid in seen:
            raise GroupUpdateError("duplicate-id", f"record ID {rid} appears twice")
        seen.add(rid)
        if type(action) is not str or action not in ACTIONS:
            raise GroupUpdateError("invalid-action", f"record {rid}: the action must be A or B")
        if reward is None:  # the rule fixes what it would be, but the record must carry it: refuse, never fill
            raise GroupUpdateError("missing-feedback", f"record {rid} has no stored score; under the rule {action} scores {REWARD[action]}, but a corrected record must be supplied")
        if not is_number(reward):
            raise GroupUpdateError("invalid-feedback-type", f"record {rid}: a score must be a number")
        if isinstance(reward, float) and not math.isfinite(reward):
            raise GroupUpdateError("nonfinite-feedback", f"record {rid}: a score must be finite")
        if reward != REWARD[action]:
            raise GroupUpdateError("contradictory-feedback", f"record {rid}: under the rule, {action} scores {REWARD[action]}")
        valid.append((rid, action, float(reward)))
    return valid


def check_scores(scores):
    """For the normalization comparisons only: 2 to 4 scores, each exactly 0 or 1."""
    if not isinstance(scores, list) or not 2 <= len(scores) <= 4:
        raise GroupUpdateError("malformed-scores", "a score list has 2 to 4 entries")
    for i, s in enumerate(scores, start=1):
        if not is_number(s) or s not in (0, 1):
            raise GroupUpdateError("invalid-score", f"score {i} must be exactly 0 or 1")
    return [float(s) for s in scores]


def normalize(scores):
    """Internal. Center on the group mean b, then divide by sigma with divisor n (here, division by four)."""
    R = np.array(scores, dtype=np.float64)                 # shape (n,)
    b = R.sum() / R.size
    sigma = np.sqrt(((R - b) ** 2).sum() / R.size)         # R.std(ddof=0); torch.std's default would divide by n - 1
    if sigma == 0:                                         # (R - b) / sigma is 0 / 0: zero weights by convention
        return b, sigma, np.zeros_like(R)
    return b, sigma, (R - b) / sigma


def sample_std_plus_epsilon(scores):
    """The inspected TRL v0.24.0 arithmetic, for comparison: divisor n - 1, plus 1e-4. Not TRL itself."""
    R = np.array(check_scores(scores), dtype=np.float64)
    c = R - R.sum() / R.size
    return c / (np.sqrt((c ** 2).sum() / (R.size - 1)) + 1e-4)


def branch(rho, w):
    """Internal. For w > 0 the term is w * min(rho, UPPER); for w < 0 it is w * max(rho, LOWER).

    So a weighted term has one kink, at the bound its weight favors, and passes smoothly through the other.
    """
    if w == 0:
        return "no-weight"                                 # zero at every policy, smooth everywhere
    if rho == (UPPER if w > 0 else LOWER):
        return "edge"                                      # the favored bound: a kink, no unique derivative
    if (rho > UPPER) if w > 0 else (rho < LOWER):
        return "clipped"                                   # past the favored bound: constant, no incentive
    if (rho < LOWER) if w > 0 else (rho > UPPER):
        return "unclipped"                                 # past the unfavorable bound: the full rho * w stays
    return "inside"                                        # includes the unfavorable bound itself


def surrogate(recs, w, pi):
    """Internal. L(pi) = mean of min(rho w, clip(rho) w), rho over the OLD probability of the SAME recorded action."""
    idx = np.array([ACTIONS.index(a) for _, a, _ in recs])  # int64, used only to index
    rho = pi[idx] / PI_OLD[idx]                              # (4,)
    clipped = np.clip(rho, LOWER, UPPER)
    terms = np.minimum(rho * w, clipped * w)
    branches = [branch(r, x) for r, x in zip(rho, w)]
    grad = None
    if "edge" not in branches:                               # d rho / dz = rho (e_a - pi) where the term moves
        active = np.array([x in ("inside", "unclipped") for x in branches])
        grad = ((active * w * rho)[:, None] * (np.eye(2)[idx] - pi)).sum(axis=0) / len(recs)
    return {
        "rho": rho, "terms": terms, "branches": branches, "L": terms.sum() / len(recs),
        "clipped_only": (clipped * w).sum() / len(recs), "unclipped": (rho * w).sum() / len(recs), "grad": grad,
        "J": float(pi @ np.array([REWARD[a] for a in ACTIONS], dtype=np.float64)),
        "kl": float((PI_OLD * (np.log(PI_OLD) - np.log(pi))).sum()),
    }


def analyze(records):
    """One group, one step, always from the same old policy."""
    recs = validate_group(records)
    b, sigma, w = normalize([r for _, _, r in recs])
    idx = np.array([ACTIONS.index(a) for _, a, _ in recs])
    scores = np.eye(2)[idx] - PI_OLD                          # grad log pi(a_i) at z_old, shape (4, 2)
    contributions = w[:, None] * scores                      # (4, 2)
    g = contributions.sum(axis=0) / 4                        # (2,): at z_old every rho is 1, inside the band
    z1 = Z_OLD + ETA * g                                     # the logits move; probabilities follow by softmax
    pi1 = softmax(z1)
    return {"recs": recs, "b": b, "sigma": sigma, "w": w, "old_prob": PI_OLD[idx], "scores": scores,
            "contributions": contributions, "g": g, "z1": z1, "pi1": pi1,
            "old": surrogate(recs, w, PI_OLD), "new": surrogate(recs, w, pi1)}


def evaluate_candidate(records, p):
    """The fixed group's L at a supplied pi = (p, 1 - p), 0.05 <= p <= 0.95. A comparison, not a step."""
    recs = validate_group(records)
    if not is_number(p) or not 0.05 <= p <= 0.95:
        raise GroupUpdateError("invalid-candidate", "pi(A) must be a number from 0.05 to 0.95")
    return surrogate(recs, normalize([r for _, _, r in recs])[2], np.array([float(p), 1 - float(p)]))


def group(order):
    return [{"id": f"e{i}", "action": a, "reward": REWARD[a]} for i, a in enumerate(order, start=1)]


def v(x, digits=6):
    return tuple(round(float(t), digits) for t in x)


# 1. The default group AABB: two As scored 1, two Bs scored 0.
AABB = analyze(group("AABB"))
print("b =", AABB["b"], "sigma =", AABB["sigma"], "(divide by 4); weights", v(AABB["w"]))
print("shapes: scores", AABB["scores"].shape, "contributions", AABB["contributions"].shape, AABB["g"].dtype)
print("contributions:", [v(c) for c in AABB["contributions"]], "-> g =", v(AABB["g"]))
print("z1 =", v(AABB["z1"]), "-> pi1 =", v(AABB["pi1"]), "; adding ETA * g to pi instead would give", v(PI_OLD + ETA * AABB["g"]))
assert AABB["b"] == 0.5 and AABB["sigma"] == 0.5 == np.std([1.0, 1.0, 0.0, 0.0]) and v(AABB["w"]) == (1, 1, -1, -1)
assert v(AABB["g"]) == (0.5, -0.5) and v(AABB["z1"]) == (0.25, -0.25)
assert math.isclose(AABB["pi1"][0], 1 / (1 + math.exp(-0.5)), rel_tol=1e-15)
old, new = AABB["old"], AABB["new"]
print("J:", old["J"], "->", round(new["J"], 12), "while the recorded mean stays", AABB["b"])
print("L at z_old:", old["L"], "gradient", v(old["grad"]), "| L at z1:", round(new["L"], 12), new["branches"])
print("ratios at z1:", v(new["rho"]), "outside [0.8, 1.2]: nothing projects them back; gradient", v(new["grad"]))
assert old["L"] == 0 and np.array_equal(old["grad"], AABB["g"]) and math.isclose(new["L"], 0.2, rel_tol=1e-15)
assert new["branches"] == ["clipped"] * 4 and not new["grad"].any() and new["L"] <= new["unclipped"]

# Central differences through the fixed-record surrogate and through J confirm both gradients at z_old.
h, w, recs = 1e-6, AABB["w"], AABB["recs"]
for c in range(2):
    up, down = Z_OLD.copy(), Z_OLD.copy()
    up[c], down[c] = h, -h
    dL = (surrogate(recs, w, softmax(up))["L"] - surrogate(recs, w, softmax(down))["L"]) / (2 * h)
    dJ = (softmax(up)[0] - softmax(down)[0]) / (2 * h)
    assert abs(dL - AABB["g"][c]) < 1e-9 and abs(dJ - (0.25, -0.25)[c]) < 1e-9

# 2. AAAB: the group mean is 0.75, but every denominator is still pi_old = 0.5.
AAAB = analyze(group("AAAB"))
print("AAAB: b =", AAAB["b"], "; old denominators", v(AAAB["old_prob"]), "; weights", v(AAAB["w"]), "; g =", v(AAAB["g"]))
assert np.array_equal(AAAB["old_prob"], [0.5] * 4)                                # not the group's 3/4 and 1/4
assert np.allclose(AAAB["w"], [3 ** -0.5] * 3 + [-(3 ** 0.5)], rtol=1e-15, atol=0)
assert np.allclose(analyze(group("ABBB"))["g"], AAAB["g"], rtol=1e-15, atol=0)   # the same step, from mean 0.25

# 3. A flat group: every score known and equal. Zero signal by convention, not an optimum.
for order in ("AAAA", "BBBB"):
    flat = analyze(group(order))
    assert flat["sigma"] == 0 and not flat["w"].any() and not flat["g"].any() and np.array_equal(flat["pi1"], PI_OLD)
    print(f"{order}: b = {flat['b']}, sigma = 0, zero weights, no step; grad J at z_old is still (0.25, -0.25)")

# 4. The fixed AABB group at supplied candidates: L(p) = min(2p - 1, 0.2).
for p in (0.6, 0.9, 0.3, 0.1):
    e = evaluate_candidate(group("AABB"), p)
    assert math.isclose(e["L"], min(2 * p - 1, 0.2), rel_tol=1e-12)
    print(f"p = {p}: L {e['L']:.6f}, clip only {e['clipped_only']:.6f}, KL(old || p) {e['kl']:.12f}, J {e['J']}")
assert evaluate_candidate(group("AABB"), 0.6)["grad"] is None                        # the favored bounds: a kink

# At p = 0.4 the A rows (w > 0) sit at 0.8 and the B rows (w < 0) at 1.2: their unfavorable bounds. The kept
# minimum is smooth there, dL/dp = 2, so the logit gradient is 2 p (1 - p) (1, -1) = (0.48, -0.48). At 0.6 the
# one-sided slopes are 2 and 0.
L_at = lambda p: evaluate_candidate(group("AABB"), p)["L"]
at_04, h = evaluate_candidate(group("AABB"), 0.4), 1e-6
assert at_04["branches"] == ["inside"] * 4 and np.allclose(at_04["grad"], [0.48, -0.48], rtol=1e-14, atol=0)
assert abs((L_at(0.4 + h) - L_at(0.4 - h)) / (2 * h) * 0.4 * 0.6 - at_04["grad"][0]) < 1e-8
left, right = (L_at(0.6) - L_at(0.6 - h)) / h, (L_at(0.6 + h) - L_at(0.6)) / h
assert abs(left - 2) < 1e-6 and abs(right) < 1e-6
print("p = 0.4: ratios", v(at_04["rho"]), "on their unfavorable bounds, smooth; gradient", v(at_04["grad"]),
      "| p = 0.6: one-sided slopes", round(left, 6), "and", round(right, 6))

# 5. All sixteen ordered groups, each with probability 1/16.
orders = ["".join(o) for o in product("AB", repeat=4)]
g_A = [analyze(group(o))["g"][0] for o in orders]
counts = [sum(o.count("A") == k for o in orders) for k in range(5)]
assert counts == [math.comb(4, k) for k in range(5)] == [1, 4, 6, 4, 1]
assert all(math.isclose(x, math.sqrt(o.count("A") * (4 - o.count("A"))) / 4, abs_tol=1e-15) for x, o in zip(g_A, orders))
E = sum(g_A) / 16


def exact_A(order, centered):
    """(1/4) sum of weight * (e_a - pi_old)_A in exact fractions, with weight R - b or the raw R."""
    R = [Fraction(REWARD[a]) for a in order]
    b = sum(R) / 4 if centered else 0
    return sum((r - b) * (Fraction(1, 2) if a == "A" else Fraction(-1, 2)) for r, a in zip(R, order)) / 4


centered, raw = (sum(exact_A(o, c) for o in orders) / 16 for c in (True, False))
print("counts", counts, "; E[g_A] =", round(E, 12), "; centering only", centered, "; raw rewards", raw, "= grad J_A")
assert math.isclose(E, 3 / 16 + math.sqrt(3) / 8, rel_tol=1e-15) and (centered, raw) == (Fraction(3, 16), Fraction(1, 4))

# 6. A different divisor, for comparison only.
print("sample std + 1e-4 on AABB:", v(sample_std_plus_epsilon([1, 1, 0, 0]), 8))


# 7. Refusals: an incomplete, invalid or contradictory record is never filled, dropped or read as a flat group.
def refusal(f, *args):
    try:
        f(*args)
    except GroupUpdateError as error:
        return error.code
    return None


missing = group("AABB")
missing[3]["reward"] = None
assert refusal(analyze, missing) == "missing-feedback"
assert refusal(analyze, group("AABB")[:3]) == "malformed-group"
assert refusal(analyze, [dict(r, reward=float("nan")) if r["id"] == "e2" else r for r in group("AABB")]) == "nonfinite-feedback"
assert refusal(analyze, [dict(r, reward=1) for r in group("AABB")]) == "contradictory-feedback"
assert refusal(analyze, [dict(r, reward=True) if r["id"] == "e1" else r for r in group("AABB")]) == "invalid-feedback-type"
assert refusal(evaluate_candidate, group("AABB"), 0.99) == "invalid-candidate"
print("refusals checked")

The helper refuses rather than repairs: no zero fill, dropped record, epsilon or clamping. Each line is its actual response. A refusal is not always ignorance: the first record’s value is fixed by the rule, but the record itself is incomplete.

• e4 chose B; its stored score is nullmissing-feedback: record e4 has no stored score. Under the stated rule action B scores 0, but a record must carry its own score: supply a corrected record; nothing is filled in
• e2 chose A; its score is NaNnonfinite-feedback: record e2: a score must be finite
• e4 chose B and is scored 1 under the rule R(B) = 0contradictory-feedback: record e4: under the stated rule, action B scores 0, not 1
• e1 scored "1", a stringinvalid-feedback-type: record e1: a score must be a number
• three records instead of fourmalformed-group: a group must have 4 entries
• two records both named e1duplicate-id: record ID e1 appears twice
• a candidate π(A) = 0.99invalid-candidate: π(A) must be a finite number from 0.05 to 0.95

Rehearsal with the worked values in view. Opening an answer is help, and nothing is recorded. This is practice, not an assessment.

1. In AABB, e3 chose B and scored 0. Why does its row push the policy toward A?

   Show an answer

   Its weight is −1 (below the mean 0.5). Its score vector is (−½, ½); times −1 it is (½, −½), the same direction as the A rows.

2. In AAAB the mean reward is 0.75. What goes in each ratio’s denominator?

   Show an answer

   π_old of the record’s own action, ½ for every record. The group’s frequencies ¾ and ¼ are not the old policy.

3. After the AABB step, L = 0.2 and the ratios are 1.2449 and 0.7551. Did clipping stop the step at the band?

   Show an answer

   No. The step was computed at π_old, where the clip is inactive, and nothing projects it back. Clipping only makes L flat once a ratio has passed the band in its favored direction.

4. AAAA, and part 5’s black-box group whose fourth score never came back, both leave π where it was. What is different?

   Show an answer

   AAAA is fully scored: zero weights by the convention used here, no relative signal. In the black-box group no rule determines the missing score, so b, σ and every weight are unknown: the calculation is unavailable, not zero. Under this page’s rule a record missing its score is also refused, but there its value is determined; the record must be corrected, not filled in silently.