Chapter 123 · Solutions
Deep RL & Applications · Solutions
Five challenges, each verified in code.
Chapter 123 · Practice Challenges, Solved¶
Five deep-RL exercises on the GridWorld and the RLHF preference data, worked with torch and scikit-learn.
In [1]:
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/"
In [2]:
import torch, torch.nn as nn, random
torch.manual_seed(0); random.seed(0); rng = np.random.default_rng(0)
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)
CHALLENGE 1
Train a DQN
Train a neural Q-network with experience replay and a target network; report its evaluation return.
In [3]:
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); buf=[]; gamma=0.95
for ep in range(400):
s=start; eps=max(0.05,1-ep/250)
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); buf.append((s,a,r,s2,done)); s=s2
if len(buf)>2000: buf.pop(0)
if len(buf)>=64:
b=random.sample(buf,64)
ss=torch.stack([feat(x[0]) for x in b]); aa=torch.tensor([x[1] for x in b])
rr=torch.tensor([x[2] for x in b],dtype=torch.float32); s2s=torch.stack([feat(x[3]) for x in b])
dd=torch.tensor([float(x[4]) for x in b])
qsa=q_net(ss).gather(1,aa[:,None]).squeeze(1)
with torch.no_grad(): tv=rr+gamma*target(s2s).max(1).values*(1-dd)
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())
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)
print('DQN evaluation return = %.2f' % evaluate(q_net))
DQN evaluation return = 9.30
Solution. The neural Q-network reaches the same near-optimal return as the tabular agent (~9.3), but it learns a function of the state, so it can scale to state spaces far too large for any table.
CHALLENGE 2
Policy gradient
Train a REINFORCE policy with a baseline; report its final return.
In [4]:
torch.manual_seed(0)
policy=nn.Sequential(nn.Linear(2,128),nn.ReLU(),nn.Linear(128,4)); opt=torch.optim.Adam(policy.parameters(),lr=5e-3)
pg=[]; baseline=0.0
for ep in range(1500):
s=start; logps=[]; rews=[]
for t in range(80):
d=torch.distributions.Categorical(logits=policy(feat(s))); a=d.sample()
logps.append(d.log_prob(a)); s,r,done=step(s,int(a)); rews.append(r)
if done: break
G=0; rets=[]
for r in reversed(rews): G=r+0.97*G; rets.insert(0,G)
rets=torch.tensor(rets); adv=rets-baseline; baseline=0.95*baseline+0.05*float(rets[0])
loss=-(torch.stack(logps)*adv).sum(); opt.zero_grad(); loss.backward(); opt.step(); pg.append(sum(rews))
print('policy-gradient final return (last 50 eps) = %.2f' % np.mean(pg[-50:]))
policy-gradient final return (last 50 eps) = 9.23
Solution. REINFORCE also solves the maze but learns more noisily than DQN, policy gradients are higher-variance, which is why a baseline (a simple critic) and, ultimately, actor-critic methods are used in practice.
CHALLENGE 3
Value-based vs policy-based
State the key difference between DQN and policy-gradient methods.
In [5]:
print('DQN (value-based) : learns Q(s,a), then acts greedily; sample-efficient, needs discrete actions')
print('REINFORCE (policy-based): learns the policy directly; handles continuous/stochastic actions, higher variance')
print('actor-critic combines both: an actor (policy) guided by a critic (value) -> PPO, the engine of RLHF')
DQN (value-based) : learns Q(s,a), then acts greedily; sample-efficient, needs discrete actions REINFORCE (policy-based): learns the policy directly; handles continuous/stochastic actions, higher variance actor-critic combines both: an actor (policy) guided by a critic (value) -> PPO, the engine of RLHF
Solution. Value-based methods learn how good actions are and act greedily; policy-based methods optimize the action probabilities directly. Actor-critic fuses them, and PPO (an actor-critic) is what RLHF uses to fine-tune language models.
CHALLENGE 4
Reward model
Fit the RLHF Bradley-Terry reward model and name the feature humans weight most.
In [6]:
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
D=Xc-Xr # one row per comparison
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))]
rm=LogisticRegression(fit_intercept=False,C=10).fit(Xtr,ytr)
w=dict(zip(F, rm.coef_[0].round(2)))
print('learned reward weights:', w)
print('most valued feature:', max(w, key=w.get))
learned reward weights: {'helpfulness': np.float64(1.28), 'factuality': np.float64(1.42), 'conciseness': np.float64(0.8), 'safety': np.float64(1.75)}
most valued feature: safety
Solution. The reward model recovers that safety and factuality carry the most weight, learned only from pairwise choices, never from the hidden true weights. This is reward modeling, stage two of RLHF.
CHALLENGE 5
Does it rank correctly?
Report the fraction of held-out comparisons the reward model ranks the way humans did.
In [7]:
margin = D[idx_te] @ rm.coef_[0]
acc = (margin > 0).mean()
print('held-out pairs ranked correctly = %.2f' % acc)
held-out pairs ranked correctly = 0.81
Solution. The reward model assigns the preferred response a higher score in most held-out comparisons (the share printed above), that scalar reward is exactly what PPO then optimizes the language model to maximize, closing the RLHF loop.
Statistics, Data Science and AI: A Visual Handbook · © 2026 John Fisher