Deep Reinforcement Learning: scaling up, and aligning AI¶
Tabular Q-learning needs one table entry per state, hopeless for a video screen, a robot's cameras, or a Go board. Deep RL replaces the table with a neural network that generalizes across states. This notebook trains a Deep Q-Network (DQN) and a policy-gradient agent on our GridWorld with torch, explains actor-critic, then builds the reward model at the heart of RLHF, the technique used to align large language models, from a teaching set of preference comparisons written for the book. Library-first with torch and scikit-learn.
import numpy as np, pandas as pd, matplotlib.pyplot as plt
import seaborn as sns # seaborn = high-level statistical plots (heatmaps, pairplots, count/bar plots)
from matplotlib.colors import ListedColormap
EM="#4338ca"; DEEP="#3730a3"; LIGHT="#c7d2fe"; INK="#1a2138"; GRID="#e6e9f2"; RED="#ef4444"; AMBER="#d97706"; GREEN="#059669"; BLUE="#2563eb"; PUR="#9333ea"; GREY="#94a3b8"; SLATE="#475569"; ORG="#4338ca"; CYAN="#0891b2"
plt.rcParams.update({"figure.facecolor":"white","axes.facecolor":"white","figure.dpi":110,"font.size":11,
"axes.edgecolor":GRID,"axes.grid":True,"grid.color":GRID,"axes.axisbelow":True,"axes.spines.top":False,
"axes.spines.right":False,"axes.titlesize":12,"axes.titleweight":"bold","legend.frameon":False})
sns.set_style("whitegrid")
BASE="https://raw.githubusercontent.com/johnfisher-ai/Statistics-Data-Science-AI-Visual-Book/main/data/"
import torch, torch.nn as nn, random
torch.manual_seed(0); random.seed(0); rng = np.random.default_rng(0)
# ---- the GridWorld from Chapter 122 (states, actions, rewards) ----
GRID=5; start,goal,trap=(0,0),(4,4),(2,2); walls={(1,3),(2,3),(3,1)}
MOVES={0:(-1,0),1:(0,1),2:(1,0),3:(0,-1)}
def step(rc,a):
r,c=rc; dr,dc=MOVES[a]; nr,nc=r+dr,c+dc
if nr<0 or nr>=GRID or nc<0 or nc>=GRID or (nr,nc) in walls: nr,nc=r,c
if (nr,nc)==goal: return (nr,nc),10.0,True
if (nr,nc)==trap: return (nr,nc),-10.0,True
return (nr,nc),-0.1,False
def feat(s): return torch.tensor([s[0]/4.0, s[1]/4.0], dtype=torch.float32) # state -> 2 numbers a net can read
print('a table needs one row per state; a 200x200 pixel screen has more states than atoms in the universe -> we need a function, not a table')
a table needs one row per state; a 200x200 pixel screen has more states than atoms in the universe -> we need a function, not a table
def make_net(): return nn.Sequential(nn.Linear(2,64), nn.ReLU(), nn.Linear(64,64), nn.ReLU(), nn.Linear(64,4))
q_net, target = make_net(), make_net(); target.load_state_dict(q_net.state_dict())
opt = torch.optim.Adam(q_net.parameters(), lr=1e-3); buffer=[]; gamma=0.95; dqn_curve=[]
for ep in range(400):
s=start; eps=max(0.05, 1-ep/250); total=0
for t in range(60):
with torch.no_grad(): qv=q_net(feat(s))
a = rng.integers(4) if rng.random()<eps else int(qv.argmax())
s2,r,done = step(s,a); total+=r; buffer.append((s,a,r,s2,done)); s=s2
if len(buffer)>2000: buffer.pop(0)
if len(buffer)>=64:
batch=random.sample(buffer,64)
ss=torch.stack([feat(b[0]) for b in batch]); aa=torch.tensor([b[1] for b in batch])
rr=torch.tensor([b[2] for b in batch],dtype=torch.float32); s2s=torch.stack([feat(b[3]) for b in batch])
dd=torch.tensor([float(b[4]) for b in batch])
qsa=q_net(ss).gather(1,aa[:,None]).squeeze(1)
with torch.no_grad(): tv=rr+gamma*target(s2s).max(1).values*(1-dd) # target network stabilizes the goal
loss=((qsa-tv)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
if done: break
if ep%20==0: target.load_state_dict(q_net.state_dict())
dqn_curve.append(total)
def evaluate(net):
outs=[]
for _ in range(50):
s=start; tot=0
for _ in range(60):
with torch.no_grad(): a=int(net(feat(s)).argmax())
s,r,done=step(s,a); tot+=r
if done: break
outs.append(tot)
return np.mean(outs)
fig, ax = plt.subplots(figsize=(8,4.2))
ax.plot(pd.Series(dqn_curve).rolling(20).mean(), color=EM, lw=2.4)
ax.set(xlabel='training episode', ylabel='total reward', title='DQN learns the maze with a neural Q-function'); plt.tight_layout(); plt.show()
print('DQN greedy evaluation return = %.2f (matches the tabular agent, but now it GENERALIZES across states)' % evaluate(q_net))
DQN greedy evaluation return = 9.30 (matches the tabular agent, but now it GENERALIZES across states)
Two tricks make it work. A plain neural net trained on RL data is unstable, so DQN adds two stabilizers. Experience replay stores past transitions in a buffer and trains on random minibatches, breaking the correlation between consecutive steps. A slowly-updated target network provides a fixed goal for the update so the net is not chasing its own moving predictions. With both, the network reaches the same near-optimal return as the tabular agent, but because it learns a function of the state, it can handle enormous or continuous state spaces a table never could. This is the algorithm that first played Atari from raw pixels.
torch.manual_seed(0)
policy = nn.Sequential(nn.Linear(2,128), nn.ReLU(), nn.Linear(128,4))
opt2 = torch.optim.Adam(policy.parameters(), lr=5e-3); pg_curve=[]; baseline=0.0
for ep in range(1500):
s=start; logps=[]; rews=[]
for t in range(80):
dist=torch.distributions.Categorical(logits=policy(feat(s))); a=dist.sample()
logps.append(dist.log_prob(a)); s,r,done=step(s,int(a)); rews.append(r)
if done: break
G=0; returns=[]
for r in reversed(rews): G=r+0.97*G; returns.insert(0,G)
returns=torch.tensor(returns); adv=returns-baseline # subtract a baseline (a simple 'critic')
baseline=0.95*baseline+0.05*float(returns[0])
loss=-(torch.stack(logps)*adv).sum(); opt2.zero_grad(); loss.backward(); opt2.step()
pg_curve.append(sum(rews))
fig, ax = plt.subplots(figsize=(8,4.2))
ax.plot(pd.Series(dqn_curve).rolling(20).mean(), color=EM, lw=2.2, label='DQN (value-based)')
ax.plot(pd.Series(pg_curve).rolling(40).mean(), color=AMBER, lw=2.2, label='policy gradient (policy-based)')
ax.set(xlabel='training episode', ylabel='total reward', title='Two families of deep RL, same GridWorld'); ax.legend(); plt.tight_layout(); plt.show()
print('policy-gradient final return (last 50 eps) = %.2f' % np.mean(pg_curve[-50:]))
policy-gradient final return (last 50 eps) = 9.23
Value-based vs policy-based. DQN learned values and acted greedily; a policy gradient method (REINFORCE) skips values and adjusts the action probabilities directly, nudging up the actions that led to high return. Its strength is handling continuous actions and stochastic policies (vital in robotics); its weakness is high variance, notice the policy-gradient curve is noisier and needs more episodes. The fix is a baseline subtracted from the return (used above), which shrinks the variance without adding bias. That baseline is really a value estimate, which leads straight to the next idea.
Combining the two. Actor-critic methods run a policy network (the actor) and a value network (the critic) together. The actor chooses actions; the critic estimates how good the current state is. A learned value used only as a baseline gives REINFORCE with a baseline. Actor-critic goes further and has the critic bootstrap, scoring each step by its reward plus the critic's estimate of the next state, so the actor gets a lower-variance learning signal at every step (at the price of some bias) instead of waiting for the episode to end. Modern algorithms, A2C, PPO, SAC, are all actor-critic variants, and PPO in particular is the workhorse behind aligning language models, which we build toward next.
print('Deep RL family tree:')
for line in [' value-based : DQN, Double-DQN, Rainbow (learn Q, act greedily)',
' policy-based : REINFORCE (learn the policy directly)',
' actor-critic : A2C, PPO, SAC (actor + critic together) <- PPO powers RLHF']:
print(line)
Deep RL family tree: value-based : DQN, Double-DQN, Rainbow (learn Q, act greedily) policy-based : REINFORCE (learn the policy directly) actor-critic : A2C, PPO, SAC (actor + critic together) <- PPO powers RLHF
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
try: pref = pd.read_excel('../../data/deep-rl-and-applications--preferences.xlsx', sheet_name='Data')
except FileNotFoundError: pref = pd.read_excel(BASE + 'deep-rl-and-applications--preferences.xlsx', sheet_name='Data')
F = ['helpfulness','factuality','conciseness','safety']
Xc = pref[[f'chosen_{f}' for f in F]].values; Xr = pref[[f'rejected_{f}' for f in F]].values
# Bradley-Terry reward model: learn weights so reward(chosen) > reward(rejected)
D = Xc - Xr # one row per comparison: chosen minus rejected
idx_tr, idx_te = train_test_split(np.arange(len(D)), test_size=0.25, random_state=0) # hold out whole comparisons
Xtr = np.vstack([D[idx_tr], -D[idx_tr]]); ytr = np.r_[np.ones(len(idx_tr)), np.zeros(len(idx_tr))] # both orders, training only
reward_model = LogisticRegression(fit_intercept=False, C=10).fit(Xtr, ytr)
w = reward_model.coef_[0]
acc = ((D[idx_te] @ w) > 0).mean()
print(f'held-out comparisons ranked correctly = {acc:.2f} ({len(idx_te)} comparisons the model never saw)')
fig, ax = plt.subplots(figsize=(7.6,4.2))
ax.bar(F, w, color=[EM if v==max(w) else LIGHT for v in w]); ax.set(ylabel='learned reward weight', title='What the reward model learned the raters value')
plt.xticks(rotation=15); plt.tight_layout(); plt.show()
print('learned priorities:', dict(zip(F, w.round(2))))
print('-> safety and factuality carry the most weight')
held-out comparisons ranked correctly = 0.81 (150 comparisons the model never saw)
learned priorities: {'helpfulness': np.float64(1.28), 'factuality': np.float64(1.42), 'conciseness': np.float64(0.8), 'safety': np.float64(1.75)}
-> safety and factuality carry the most weight
How RLHF works, and what we just did. Aligning a language model has three stages. (1) Supervised fine-tuning on good example answers. (2) Reward modeling: humans compare pairs of responses, and a model learns to predict their preference, which is exactly the Bradley-Terry logistic model we just fit, learning that safety and factuality matter most from nothing but pairwise choices (it ranks held-out comparisons correctly at the rate printed above). (3) RL fine-tuning: a policy-gradient method (PPO, an actor-critic) updates the language model to maximize that learned reward, while a penalty keeps it close to the original so it does not game the reward. The reward model is the bridge from messy human preference to a trainable signal, and it is a plain classifier at heart.
chosen_margin = D[idx_te] @ w # reward(chosen) - reward(rejected) on the held-out comparisons
fig, ax = plt.subplots(figsize=(7.6,4.0))
ax.hist(chosen_margin, bins=25, color=GREEN, alpha=0.8)
ax.axvline(0, color=RED, ls='--', lw=2)
ax.set(xlabel='reward(chosen) - reward(rejected)', ylabel='count', title='The reward model scores the preferred responses higher (mostly > 0)')
plt.tight_layout(); plt.show()
print(f'{(chosen_margin > 0).sum()} of {len(chosen_margin)} held-out comparisons have a positive margin')
122 of 150 held-out comparisons have a positive margin
Reading the check. For held-out comparisons the reward model assigns the preferred response a higher reward most of the time (the margin is mostly positive). That single number, a scalar reward the model can compute for any response, is what PPO then optimizes. Everything in this chapter converges here: a neural network (Demo 1) trained by a policy-gradient actor-critic (Demos 2 to 3) against a reward learned from preference comparisons (Demo 4), which in a real system come from human raters. That pipeline is a large part of why modern chatbots are helpful and safe.
Deep RL and its applications, in one view¶
- Deep RL swaps the Q-table for a neural network so RL scales to huge or continuous state spaces.
- DQN (value-based) uses experience replay and a target network to learn a stable Q-function; it reached human-level play on many Atari games from pixels.
- Policy-gradient methods learn the policy directly (great for continuous control) but are high-variance; a baseline/critic tames it.
- Actor-critic (A2C, PPO, SAC) combines both; PPO is the engine of RLHF.
- RLHF aligns LLMs: humans rank responses -> a reward model (a classifier) learns their preference -> PPO optimizes the model against it.