diff --git a/AI-Engineering-Preparation/labs/02-rl-from-scratch/preferences.py b/AI-Engineering-Preparation/labs/02-rl-from-scratch/preferences.py index 46c6153..1dec0a8 100644 --- a/AI-Engineering-Preparation/labs/02-rl-from-scratch/preferences.py +++ b/AI-Engineering-Preparation/labs/02-rl-from-scratch/preferences.py @@ -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}) @@ -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: @@ -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:])) diff --git a/AI-Engineering-Preparation/labs/02-rl-from-scratch/solution/lab.ipynb b/AI-Engineering-Preparation/labs/02-rl-from-scratch/solution/lab.ipynb index 26784bf..bd1a960 100644 --- a/AI-Engineering-Preparation/labs/02-rl-from-scratch/solution/lab.ipynb +++ b/AI-Engineering-Preparation/labs/02-rl-from-scratch/solution/lab.ipynb @@ -4,107 +4,29 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Lab 02 (solution): RL from scratch\n", + "# Solution: preference learning\n", "\n", - "Reference implementation. Two learning signals built by hand: value-based RL (Q-learning on a gridworld) and preference-based reward modeling (the Bradley-Terry core of RLHF). Concepts: [rl-primitives](../../../concepts/rl/rl-primitives.md), [rlhf](../../../concepts/rl/rlhf.md)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Step 0: Q-learning on a gridworld" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "execution_count": null, - "outputs": [], - "source": [ - "import sys\n", - "sys.path.insert(0, \"..\") # modules (qlearning.py, preferences.py) live in the lab root\n", - "from qlearning import train, greedy_path, render_policy, OPTIMAL_STEPS\n", - "# RL has three parts: an environment (the grid), a reward (+10 at goal, -1 per step), and a policy\n", - "# learned from values. Q-learning fills a table Q[state][action] and acts greedily on it.\n", - "Q = train(seed=0)\n", - "print(render_policy(Q))\n", - "path = greedy_path(Q)\n", - "print(f\"\\npath: {path} ({len(path)-1} steps, optimal={OPTIMAL_STEPS})\")\n" - "print(f\"\\npath: {path} ({len(path)-1} steps, optimal={OPTIMAL_STEPS})\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Step 1: value is long-run, not immediate" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "execution_count": null, - "outputs": [], - "source": [ - "# The \"value\" is long-run, not immediate: the start state has positive value because the policy from\n", - "# there eventually reaches the +10 goal, despite paying step penalties on the way.\n", - "print(\"value of best action at start:\", round(max(Q[0]), 2))\n" - "print(\"value of best action at start:\", round(max(Q[0]), 2))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Step 2: reward modeling from preferences (RLHF core)" + "This solution fits a small Bradley-Terry reward model from deterministic pairwise preferences." ] }, { "cell_type": "code", - "metadata": {}, "execution_count": null, - "outputs": [], - "source": [ - "from preferences import make_preferences, fit_reward, ranking, TRUE_QUALITY\n", - "# RLHF's engine: you can't write a reward for \"better answer,\" so you learn one from comparisons.\n", - "# Bradley-Terry: P(i preferred over j) = sigmoid(r_i - r_j). Fit r to the preferences.\n", - "prefs = make_preferences(TRUE_QUALITY) # only \"A beats B\" pairs - no scores\n", - "r = fit_reward(prefs, len(TRUE_QUALITY))\n", - "print(f\"{len(prefs)} pairwise prefs -> learned reward:\", [round(x,2) for x in r])\n", - "print(\"learned ranking:\", ranking(r), \" true ranking:\", ranking(TRUE_QUALITY))\n" - "print(\"learned ranking:\", ranking(r), \" true ranking:\", ranking(TRUE_QUALITY))" - ] - }, - { - "cell_type": "markdown", "metadata": {}, - "source": [ - "## Step 3: the policy picks the best option" - ] - }, - { - "cell_type": "code", - "metadata": {}, - "execution_count": null, "outputs": [], "source": [ - "from preferences import policy_choice\n", - "# The policy optimized against the reward model picks the top-scored option - learned with no labels,\n", - "# only preferences. This is the loop RLHF runs at scale: preferences -> reward model -> policy.\n", - "print(\"policy picks item:\", policy_choice(r), \"(true best:\", ranking(TRUE_QUALITY)[0], \")\")\n" - "print(\"policy picks item:\", policy_choice(r), \"(true best:\", ranking(TRUE_QUALITY)[0], \")\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## What you built\n", + "import sys\n", + "\n", + "sys.path.insert(0, \"..\")\n", + "\n", + "from preferences import TRUE_QUALITY, fit_reward, make_preferences, policy_choice, ranking\n", "\n", - "The two learning signals behind modern model training, from scratch. **Q-learning** is value-based RL: an agent learns Q[state][action] by trial and reward and acts greedily, converging to the optimal six-step route - environment, reward, and policy made concrete. **Reward modeling** is the engine inside RLHF: from pairwise preferences alone (no scores), a Bradley-Terry fit recovers a reward that orders the options, and a greedy policy picks the best. \n", + "preferences = make_preferences()\n", + "reward = fit_reward(preferences)\n", "\n", - "**Where this simplifies:** the gridworld is tiny and deterministic, so a table suffices; real RL uses function approximation (a network) over enormous state spaces and methods like PPO for stability. The reward model here fits noiseless preferences over five fixed items; in RLHF the items are model outputs, the preferences are noisy human labels, and a policy is then optimized against the learned reward with a KL penalty to stay close to the base model. Concepts: [rl-primitives](../../../concepts/rl/rl-primitives.md), [policy-gradients](../../../concepts/rl/policy-gradients.md), [rlhf](../../../concepts/rl/rlhf.md)." + "print(\"true ranking:\", ranking(TRUE_QUALITY))\n", + "print(\"reward ranking:\", ranking(reward))\n", + "print(\"policy choice:\", policy_choice(reward))\n" ] } ], @@ -116,9 +38,9 @@ }, "language_info": { "name": "python", - "version": "3.11" + "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 5 -} \ No newline at end of file +}