Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 52 additions & 156 deletions AI-Engineering-Preparation/labs/02-rl-from-scratch/preferences.py
Original file line number Diff line number Diff line change
@@ -1,61 +1,49 @@
#!/usr/bin/env python3
"""Reward modeling from preferences (Lab 02).

How do you train a model toward something you cannot write down a reward function for - "be more
helpful," "this answer is better"? You collect comparisons (A is better than B) and fit a reward
that explains them. This is the engine inside RLHF: a reward model learned from human pairwise
preferences, which a policy is then optimized against.

This builds the reward model from scratch with the Bradley-Terry model: the probability that item i is
preferred over j is the logistic function of their reward difference, sigma(r_i - r_j). Fitting r by
gradient ascent on the preference log-likelihood recovers a scoring that orders items the same way the
preferences do - and a greedy policy then picks the top-scored item. No labeled scores are ever shown
to the learner, only which of two was preferred. Deterministic, offline, standard-library only.

References: Bradley & Terry (1952); Christiano et al. (2017), Deep RL from Human Preferences,
arXiv:1706.03741; Ouyang et al. (2022), InstructGPT, arXiv:2203.02155.

Usage:
python preferences.py --self-test
python preferences.py --demo
"""
"""Tiny preference-learning demo for RLHF-style reward modeling."""

from __future__ import annotations

import argparse
import math
import sys

TRUE_QUALITY = {
"brief-safe": 0.35,
"detailed-helpful": 0.85,
"verbose-unclear": 0.20,
"concise-grounded": 0.95,
"polite-vague": 0.45,
TRUE_QUALITY: dict[str, float] = {
"concise": 1.6,
"correct": 2.2,
"verbose": 0.4,
"unsafe": -1.4,
"vague": -0.7,
}


def sigmoid(value: float) -> float:
return 1.0 / (1.0 + math.exp(-value))
def sigmoid(x: float) -> float:
"""Numerically stable sigmoid."""
if x >= 0:
z = math.exp(-x)
return 1.0 / (1.0 + z)

z = math.exp(x)
return z / (1.0 + z)


def make_preferences() -> list[tuple[str, str]]:
"""Create noiseless pairwise preferences from the hidden true quality."""
"""Create deterministic pairwise preferences from hidden quality scores."""
items = list(TRUE_QUALITY)
pairs = []
preferences: list[tuple[str, str]] = []

for i, left in enumerate(items):
for right in items[i + 1 :]:
if TRUE_QUALITY[left] >= TRUE_QUALITY[right]:
pairs.append((left, right))
preferences.append((left, right))
else:
pairs.append((right, left))
preferences.append((right, left))

return pairs
return preferences


def fit_reward(
preferences: list[tuple[str, str]],
epochs: int = 500,
lr: float = 0.08,
lr: float = 0.1,
) -> dict[str, float]:
"""Fit a Bradley-Terry reward model from pairwise preferences."""
items = sorted({item for pair in preferences for item in pair})
Expand All @@ -65,14 +53,13 @@ def fit_reward(
grad = dict.fromkeys(items, 0.0)

for winner, loser in preferences:
margin = reward[winner] - reward[loser]
probability = sigmoid(margin)
error = 1.0 - probability
grad[winner] += error
grad[loser] -= error
probability = sigmoid(reward[winner] - reward[loser])
update = 1.0 - probability
grad[winner] += update
grad[loser] -= update

for item in items:
reward[item] += lr * grad[item]
reward[item] += lr * grad[item] / max(1, len(preferences))

mean_reward = sum(reward.values()) / len(reward)
for item in items:
Expand All @@ -87,142 +74,51 @@ def ranking(scores: dict[str, float]) -> list[str]:


def policy_choice(reward: dict[str, float]) -> str:
"""A policy optimized against the reward model chooses the top reward item."""
"""Choose the highest-reward item."""
return ranking(reward)[0]


def _self_test() -> int:
prefs = make_preferences()
reward = fit_reward(prefs)
learned = ranking(reward)
truth = ranking(TRUE_QUALITY)
def _self_test() -> None:
preferences = make_preferences()
reward = fit_reward(preferences)

agree = 0
for winner, loser in prefs:
if reward[winner] > reward[loser]:
agree += 1
assert len(preferences) == 10
assert policy_choice(reward) == ranking(TRUE_QUALITY)[0]
assert ranking(reward)[-1] == ranking(TRUE_QUALITY)[-1]

assert learned == truth, (learned, truth)
assert policy_choice(reward) == truth[0], (policy_choice(reward), truth[0])
assert agree == len(prefs), (agree, len(prefs))
print("preferences self-test passed")

print(
"self-test: deterministic Bradley-Terry reward model recovered "
f"the true ranking {learned} from {len(prefs)} pairwise preferences; "
f"policy picks item {policy_choice(reward)}; "
f"{agree}/{len(prefs)} preferences satisfied"
)

return 0
def _demo() -> None:
preferences = make_preferences()
reward = fit_reward(preferences)

print("Preferences:")
for winner, loser in preferences:
print(f" {winner} > {loser}")

def _demo() -> int:
prefs = make_preferences()
reward = fit_reward(prefs)

print("Learned reward scores:")
for item in ranking(reward):
print(f"- {item:18s} {reward[item]: .3f}")

print("\nLearned reward ranking:", ranking(reward))
print("True quality ranking: ", ranking(TRUE_QUALITY))
print("Policy choice:", policy_choice(reward))
print("True best:", ranking(TRUE_QUALITY)[0])

import sys

# Latent "true" quality of five candidate answers. The learner never sees these - only comparisons.
TRUE_QUALITY = [0.1, 0.9, 0.5, 0.3, 0.7]


def make_preferences(true_quality: list[float]) -> list[tuple[int, int]]:
"""Every ordered pair (winner, loser) where the higher-quality item wins. Noiseless here so the
test is deterministic; real preference data is noisy and the same fit still works."""
n = len(true_quality)
return [(i, j) for i in range(n) for j in range(n)
if i != j and true_quality[i] > true_quality[j]]


def _sigmoid(x: float) -> float:
return 1.0 / (1.0 + math.exp(-x))


def fit_reward(prefs: list[tuple[int, int]], n: int, iters: int = 2000, lr: float = 0.1) -> list[float]:
"""Gradient ascent on the Bradley-Terry log-likelihood. For a preference (i over j), the gradient
pushes r_i up and r_j down by (1 - sigma(r_i - r_j)) - the model's surprise that i won. Rewards are
centered each step because only differences are identifiable (a constant shift changes nothing)."""
r = [0.0] * n
for _ in range(iters):
grad = [0.0] * n
for i, j in prefs:
err = 1.0 - _sigmoid(r[i] - r[j])
grad[i] += err
grad[j] -= err
for k in range(n):
r[k] += lr * grad[k] / len(prefs)
mean = sum(r) / n
r = [x - mean for x in r]
return r


def ranking(scores: list[float]) -> list[int]:
return sorted(range(len(scores)), key=lambda k: scores[k], reverse=True)


def policy_choice(scores: list[float]) -> int:
"""A reward-greedy policy: pick the highest-reward option."""
return ranking(scores)[0]


def _self_test() -> int:
prefs = make_preferences(TRUE_QUALITY)
r = fit_reward(prefs, len(TRUE_QUALITY))
assert fit_reward(prefs, len(TRUE_QUALITY)) == r # deterministic

learned = ranking(r)
truth = ranking(TRUE_QUALITY)
assert learned == truth, (learned, truth) # recovered the true order from comparisons alone

# the policy picks the genuinely best answer, learned only from "A beats B" signals
assert policy_choice(r) == truth[0]

# pairwise consistency: learned reward agrees with every training preference
agree = sum(1 for i, j in prefs if r[i] > r[j])
assert agree == len(prefs), (agree, len(prefs))

print(f"self-test: deterministic fit; recovered true ranking {learned} from {len(prefs)} pairwise "
f"prefs (no scores shown); policy picks item {policy_choice(r)}; {agree}/{len(prefs)} prefs satisfied")
return 0


def main() -> int:
parser = argparse.ArgumentParser(description="Reward modeling from pairwise preferences")
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--self-test", action="store_true")
parser.add_argument("--demo", action="store_true")
args = parser.parse_args()
args = parser.parse_args(argv)

if args.self_test:
return _self_test()
_self_test()
return 0

if args.demo:
return _demo()
_demo()
return 0

parser.print_help()
ap = argparse.ArgumentParser(description="Reward modeling from pairwise preferences")
ap.add_argument("--self-test", action="store_true")
ap.add_argument("--demo", action="store_true")
args = ap.parse_args()
if args.self_test:
return _self_test()
prefs = make_preferences(TRUE_QUALITY)
r = fit_reward(prefs, len(TRUE_QUALITY))
if args.demo:
print("learned reward (centered):", [round(x, 2) for x in r])
print("learned ranking:", ranking(r), " true ranking:", ranking(TRUE_QUALITY))
print("policy picks item:", policy_choice(r))
else:
print("learned ranking:", ranking(r))
return 0


if __name__ == "__main__":
raise SystemExit(main())
sys.exit(main())
raise SystemExit(main(sys.argv[1:]))
Loading
Loading