def run_reinforce(grid):
"""REINFORCE @eq-reinforce: one full episode in, ONE parameter update out."""
theta = np.zeros((N_S, 2))
curve = np.full(len(grid), objective(theta))
k, steps, n_updates = 1, 0, 0
while steps < BUDGET:
s, traj = 0, []
for _ in range(T_MAX): # 1. generate a full episode
p1 = prob_back(theta, s)
a = 1 if RNG.random() < p1 else 0
s2, r, done = step(s, a)
traj.append((s, a, r, p1))
s = s2
if done:
break
steps += len(traj)
# The whole episode is scored first, then theta moves once: every score below is
# evaluated at the theta that generated the episode, per @eq-reinforce. The three
# lines touching g are the softmax score of @exm-softmax written out by hand for
# two actions, and g accumulates the entire sum before theta moves at all.
g, G = np.zeros((N_S, 2)), 0.0 # 2. returns, accumulated backwards
for t in range(len(traj) - 1, -1, -1):
s_, a_, r_, p1_ = traj[t]
G = r_ + GAMMA * G # G_t, the Monte-Carlo estimate of q_pi(s_t, a_t)
d = GAMMA**t * G
g[s_, a_] += d
g[s_, 0] -= d * (1.0 - p1_)
g[s_, 1] -= d * p1_
theta += ALPHA_THETA * g # 3. ONE update, after the episode
n_updates += 1
k = record(curve, k, grid, steps, theta)
return curve, n_updates
def run_qac(grid):
"""QAC @eq-qac-critic and @eq-qac-actor: one transition in, one update out.
Not literally the same actor as REINFORCE, in two ways. The estimate of q_pi is the
critic's q_sa, read BEFORE the critic moves, because @eq-qac-actor asks for
qhat(s_t, a_t, w_t); and the summand carries no gamma**t, because QAC implements
@eq-spg-simple, whose discounting has already migrated into d_{pi,mu}. The critic is a
table q(s,a) = w[s,a], which is @eq-sarsa-fa with indicator features: grad_w qhat is
then the one-hot vector picking out (s,a), so the update touches that entry alone.
There is no traj list and no backward pass, because nothing has to be remembered until
the end, and theta moves inside the transition loop rather than after it.
"""
theta, w = np.zeros((N_S, 2)), np.zeros((N_S, 2)) # actor table, critic table
curve = np.full(len(grid), objective(theta))
k, steps = 1, 0
while steps < BUDGET:
s = 0
p1 = prob_back(theta, s)
a = 1 if RNG.random() < p1 else 0
for _ in range(T_MAX):
s2, r, done = step(s, a) # 1. interact
if done:
target = r
else:
p1n = prob_back(theta, s2)
a2 = 1 if RNG.random() < p1n else 0
target = r + GAMMA * w[s2, a2] # bootstrap, not a realized return
q_sa = w[s, a]
w[s, a] += ALPHA_W * (target - q_sa) # 2. critic, @eq-qac-critic
# The score must be read at the CURRENT theta. Reusing the p1 cached when
# this state was entered is stale whenever the actor already moved row s,
# which happens on every wall bump where s2 == s.
pp = prob_back(theta, s)
theta[s, a] += ALPHA_THETA * q_sa # 3. actor, one per transition
theta[s, 0] -= ALPHA_THETA * q_sa * (1.0 - pp)
theta[s, 1] -= ALPHA_THETA * q_sa * pp
steps += 1
k = record(curve, k, grid, steps, theta)
if done or steps >= BUDGET:
break
s, a, p1 = s2, a2, p1n
return curve # its update count is the step count, so there is nothing to report
grid = np.linspace(0, BUDGET, 201)
runs = [run_reinforce(grid) for _ in range(N_RUNS)]
rein, n_up = np.array([c for c, _ in runs]), np.mean([n for _, n in runs])
qac = np.array([run_qac(grid) for _ in range(N_RUNS)])
# ponytail: one runnable check of the figure's central claim, on the simulated runs
# themselves. Per environment step QAC's median must lead over the whole second half of the
# budget, and REINFORCE's spread across runs must be the wider one: faster AND steadier out
# of the same experience, which is what the two estimates of q_pi buy at a common step size.
m_rein, m_qac = np.median(rein, axis=0), np.median(qac, axis=0)
iqr = {k: np.percentile(a, 75, axis=0) - np.percentile(a, 25, axis=0)
for k, a in (("rein", rein), ("qac", qac))}
half = len(m_qac) // 2
assert np.all(m_qac[half:] > m_rein[half:]), "QAC median does not lead throughout"
assert np.all(iqr["rein"][[half, -1]] > iqr["qac"][[half, -1]]), "REINFORCE not the wider"
fig, ax = new_ax(6.8, 4.6)
j_opt = GAMMA ** (N_S - 1) * R_GOAL
ax.plot([0, BUDGET], [j_opt, j_opt], color=GREY, lw=0.9, ls="--", zorder=2)
for arr, colour in ((rein, COL["sky"]), (qac, COL["vermilion"])):
lo, med, hi = np.percentile(arr, [25, 50, 75], axis=0)
ax.fill_between(grid, lo, hi, color=colour, alpha=0.22, lw=0, zorder=3)
ax.plot(grid, med, color=colour, lw=2.0, zorder=5)
ax.set_xlim(0, BUDGET * 1.02)
ax.set_ylim(2.2, 8.2) # headroom above the dashed rule carries its tag, not data
ax.set_xticks([0, 5000, 10000, 15000, 20000])
ax.set_xticklabels(["0", "5k", "10k", "15k", "20k"])
ax.set_xlabel("environment steps (transitions consumed)")
ax.set_ylabel(r"$J(\theta)=v_{\pi}(s_{0})$")
# Every label sits in space no curve or band enters: the reference tag above the dashed
# rule, "QAC" in the empty wedge above its own steep rise, "REINFORCE" under its own curve,
# and the two notes in the whole empty lower-right quadrant.
label_at(ax, 300, 7.94, r"$J$ at the optimal policy", color=GREY, fontsize=8.5)
label_at(ax, 3000, 6.95, "QAC", color=COL["vermilion"], fontsize=10.5)
label_at(ax, 6300, 5.25, "REINFORCE", color=COL["sky"], fontsize=10.5)
label_at(ax, 11200, 4.55, f"shaded: interquartile range across {N_RUNS} runs",
color=GREY, fontsize=8.5, va="top")
label_at(ax, 11200, 3.85,
"the same 20k transitions buy QAC 20,000\n"
f"updates and REINFORCE {n_up:,.0f}, one per episode",
color=INK, fontsize=8.5, va="top")
ax.set_title(r"Same actor, same step size: the estimate of $q_{\pi}$ is the"
" whole difference", fontsize=11)
plt.show()
