import plotly.io as pio; pio.renderers.default = "notebook_connected"

import time

import numpy as np
import torch
import plotly.graph_objects as go
from plotly.subplots import make_subplots

NAVY, GOLD, GREEN, RED, GRAY, MUD_COLOR = '#002676', '#FDB515', '#028842', '#C0362C', '#808080', '#C8A97E'

H, W = 4, 4
GOAL, PIT, WALL, START = (0, 3), (1, 3), (1, 1), (3, 0)
MUD = [(2, 1), (2, 2)]
ACTIONS = ['up', 'right', 'down', 'left']
MOVE = {'up': (-1, 0), 'right': (0, 1), 'down': (1, 0), 'left': (0, -1)}
SIDE = {'up': ('left', 'right'), 'down': ('left', 'right'), 'left': ('up', 'down'), 'right': ('up', 'down')}
SLIP, STEP_COST, MUD_COST, T_MAX = 0.1, -0.04, -0.3, 40

STATES = [(r, c) for r in range(H) for c in range(W) if (r, c) != WALL]
IDX = {s: i for i, s in enumerate(STATES)}


def move(s, a):
    r, c = s[0] + MOVE[a][0], s[1] + MOVE[a][1]
    return s if not (0 <= r < H and 0 <= c < W) or (r, c) == WALL else (r, c)


def transition(s, a):
    """p(s' | s, a) as a dict {s': probability}: the intended move 0.8, each sideways move 0.1."""
    out = {}
    for b, p in ((a, 1 - 2 * SLIP), (SIDE[a][0], SLIP), (SIDE[a][1], SLIP)):
        s2 = move(s, b)
        out[s2] = out.get(s2, 0) + p
    return out


def reward(s, a):
    """r(s, a): +1 / -1 for the exit action at the goal / pit, -0.3 for a step in the mud, else -0.04."""
    if s == GOAL:
        return 1.0
    if s == PIT:
        return -1.0
    return MUD_COST if s in MUD else STEP_COST


print('choose up at (3, 1):  ', transition((3, 1), 'up'))
print('choose up at (3, 0):  ', transition((3, 0), 'up'), ' (sliding left hits the edge: stay put)')
print('choose right at (1, 0):', transition((1, 0), 'right'), ' (the wall: stay put)')

choose up at (3, 1):   {(2, 1): 0.8, (3, 0): 0.1, (3, 2): 0.1}
choose up at (3, 0):   {(2, 0): 0.8, (3, 0): 0.1, (3, 1): 0.1}  (sliding left hits the edge: stay put)
choose right at (1, 0): {(1, 0): 0.8, (0, 0): 0.1, (2, 0): 0.1}  (the wall: stay put)

def grid_figure(title=None, size=420, fig=None, col=None):
    """A plotly figure with the grid drawn: cells, wall, mud, goal, pit. Cell (i, j) is centered at x = j, y = i.
    Pass fig and col to draw into panel col of a figure made with make_subplots(1, n)."""
    new = fig is None
    fig = go.Figure() if new else fig
    where = {} if col is None else dict(row=1, col=col)
    for r in range(H):
        for c in range(W):
            s = (r, c)
            fill = '#404040' if s == WALL else MUD_COLOR if s in MUD else '#D5ECDD' if s == GOAL else '#F6D5CF' if s == PIT else '#F2F2F2'
            fig.add_shape(type='rect', x0=c - 0.5, x1=c + 0.5, y0=r - 0.5, y1=r + 0.5, fillcolor=fill, line=dict(color='white', width=3), layer='below', **where)
            label = {GOAL: 'goal +1', PIT: 'pit −1', WALL: 'wall', START: 'start'}.get(s, 'mud −0.3' if s in MUD else '')
            if label:
                fig.add_annotation(x=c, y=r + 0.32, text=label, showarrow=False, font=dict(size=12, color='white' if s == WALL else 'black'), **where)
    fig.update_xaxes(range=[-0.5, W - 0.5], visible=False, **where)
    fig.update_yaxes(range=[H - 0.5, -0.5], visible=False, scaleanchor='x' if col in (None, 1) else f'x{col}', **where)
    if new:
        fig.update_layout(width=size, height=size, margin=dict(l=10, r=10, t=40 if title else 10, b=10), title=title,
                          plot_bgcolor='white', showlegend=False)
    return fig


grid_figure('The grid world')

def transition_figure(s, a):
    """p(s' | s, a) as arrows from s: the thicker the arrow, the more likely the move."""
    fig = grid_figure(f'p(s\' | s = {s}, a = {a})')
    for s2, prob in transition(s, a).items():
        dy, dx = s2[0] - s[0], s2[1] - s[1]
        col = NAVY if prob > 0.5 else GRAY
        fig.add_annotation(x=s[1] + 0.62 * dx, y=s[0] + 0.62 * dy, ax=s[1] + 0.2 * dx, ay=s[0] + 0.2 * dy, xref='x', yref='y', axref='x',
                           ayref='y', showarrow=True, arrowhead=2, arrowwidth=2 + 6 * prob, arrowcolor=col, text='')
        fig.add_annotation(x=s2[1] + 0.12 * dx, y=s2[0] - 0.3, text=f'<b>{prob:.1f}</b>', showarrow=False, font=dict(size=18, color=col))
    fig.add_shape(type='circle', x0=s[1] - 0.17, x1=s[1] + 0.17, y0=s[0] - 0.17, y1=s[0] + 0.17, fillcolor=NAVY, line=dict(width=0))
    return fig


transition_figure((3, 1), 'up')

def policy(theta, s):
    """pi(. | s): softmax of the cell's row of logits."""
    z = theta[IDX[s]] - theta[IDX[s]].max()
    p = np.exp(z)
    return p / p.sum()


def grad_log_pi(theta, s, a):
    """grad of ln pi(a | s) with respect to the whole table: e_a - pi in row s, zero elsewhere."""
    g = np.zeros_like(theta)
    g[IDX[s]] = -policy(theta, s)
    g[IDX[s], a] += 1.0
    return g


theta0 = np.zeros((len(STATES), 4))          # all logits 0: the uniform policy
print('parameters:', theta0.size)
print('uniform policy at the start cell:', {a: float(p) for a, p in zip(ACTIONS, policy(theta0, START).round(2))})

parameters: 60
uniform policy at the start cell: {'up': 0.25, 'right': 0.25, 'down': 0.25, 'left': 0.25}

theta_test = np.random.default_rng(0).normal(size=theta0.shape)
t = torch.tensor(theta_test, requires_grad=True)
s, a = (2, 0), ACTIONS.index('up')
torch.log_softmax(t[IDX[s]], 0)[a].backward()
print('autograd :', t.grad[IDX[s]].numpy().round(4))
print('e_a - pi :', grad_log_pi(theta_test, s, a)[IDX[s]].round(4))

autograd : [ 0.7929 -0.408  -0.1193 -0.2656]
e_a - pi : [ 0.7929 -0.408  -0.1193 -0.2656]

def policy_figure(theta, title, fig=None, col=None):
    """Arrows whose length is the probability of each move."""
    fig = grid_figure(title, fig=fig, col=col)
    ref = '' if col in (None, 1) else str(col)
    for s in STATES:
        if s in (GOAL, PIT):
            continue
        p = policy(theta, s)
        for k, a in enumerate(ACTIONS):
            dy, dx = MOVE[a]
            L = 0.08 + 0.34 * p[k]
            fig.add_annotation(x=s[1] + L * dx, y=s[0] + L * dy, ax=s[1], ay=s[0], xref=f'x{ref}', yref=f'y{ref}', axref=f'x{ref}', ayref=f'y{ref}',
                               showarrow=True, arrowhead=2, arrowwidth=1 + 3 * p[k], arrowcolor=GOLD if p[k] == p.max() and p.max() > 0.3 else NAVY, text='')
    return fig


def policy_pair(theta_left, theta_right, titles):
    """Two policies side by side (the "Policies in the Grid World" slide)."""
    fig = make_subplots(rows=1, cols=2, subplot_titles=titles, horizontal_spacing=0.05)
    for col, theta in ((1, theta_left), (2, theta_right)):
        policy_figure(theta, None, fig=fig, col=col)
    fig.update_layout(width=820, height=440, margin=dict(l=10, r=10, t=40, b=10), plot_bgcolor='white', showlegend=False)
    return fig


policy_figure(theta0, 'Uniform policy: every action 0.25')

def rollout(theta, rng):
    """One trajectory: lists of states, actions (None for the exit action), and rewards, plus the log-probabilities of
    the policy's choices and of the world's transitions."""
    s, S, A, R = START, [], [], []
    log_pi, log_world = 0.0, 0.0                     # p(s_0) = 1: always the start cell
    for t in range(T_MAX):
        if s in (GOAL, PIT):
            S.append(s); A.append(None); R.append(reward(s, None))
            break
        p = policy(theta, s)
        a = rng.choice(4, p=p)
        S.append(s); A.append(a); R.append(reward(s, ACTIONS[a]))
        nxt = transition(s, ACTIONS[a])
        k = rng.choice(len(nxt), p=list(nxt.values()))
        log_pi += np.log(p[a])
        log_world += np.log(list(nxt.values())[k])
        s = list(nxt)[k]
    return S, A, R, log_pi, log_world


rng = np.random.default_rng(4)
S, A, R, log_pi, log_world = rollout(theta0, rng)
while len(S) > 10:                                # find a short one to print
    S, A, R, log_pi, log_world = rollout(theta0, rng)
print('states :', S)
print('actions:', [ACTIONS[a] if a is not None else 'exit' for a in A])
print('rewards:', [round(r, 2) for r in R], '  return R =', round(sum(R), 2))
print(f'ln p(tau) = ln p(s0) + sum ln pi + sum ln p(s\'|s,a) = 0 + {log_pi:.2f} + {log_world:.2f} = {log_pi + log_world:.2f}')

states : [(3, 0), (2, 0), (2, 1), (2, 2), (3, 2), (2, 2), (1, 2), (1, 3)]
actions: ['up', 'right', 'right', 'down', 'up', 'up', 'right', 'exit']
rewards: [-0.04, -0.04, -0.3, -0.3, -0.04, -0.3, -0.04, -1.0]   return R = -2.06
ln p(tau) = ln p(s0) + sum ln pi + sum ln p(s'|s,a) = 0 + -9.70 + -1.56 = -11.27

def estimate_J(theta, n=2000, seed=0):
    rng = np.random.default_rng(seed)
    returns = [sum(rollout(theta, rng)[2]) for _ in range(n)]
    return np.mean(returns), np.std(returns) / np.sqrt(n)

J, se = estimate_J(theta0)
print(f'uniform policy: J ≈ {J:.2f} ± {2 * se:.2f}  (it wanders about {np.mean([len(rollout(theta0, rng)[0]) for _ in range(500)]):.0f} steps, paying the step cost)')

uniform policy: J ≈ -2.61 ± 0.06  (it wanders about 29 steps, paying the step cost)

def trajectory_figure(theta, trajs, title):
    fig = grid_figure(title, size=460)
    for (S, A, R, *_), col, off in zip(trajs, [GREEN, GOLD, RED], [-0.13, 0.0, 0.13]):
        fig.add_trace(go.Scatter(x=[s[1] + off for s in S], y=[s[0] + off for s in S], mode='lines+markers',
                                 line=dict(color=col, width=4), name=f'R = {sum(R):+.2f} ({len(S)} steps)'))
    fig.update_layout(showlegend=True, legend=dict(orientation='h', y=-0.02))
    return fig


rng = np.random.default_rng(2)
picked = {}                                      # short, non-looping examples of three outcomes (most are much longer)
while len(picked) < 3:
    tr = rollout(theta0, rng)
    S = tr[0]
    if len(S) <= 12 and len(set(S)) == len(S):
        kind = 'pit' if S[-1] == PIT else 'mud' if set(S) & set(MUD) else 'clean'
        picked.setdefault(kind, tr)
trajectory_figure(theta0, [picked['clean'], picked['mud'], picked['pit']], 'Three short trajectories of the uniform policy')

def grad_estimate(theta, traj, reward_to_go=False):
    """One trajectory's REINFORCE gradient: sum_t grad ln pi(a_t | s_t) * weight_t, where the weight is the whole return
    R(tau) or, with reward_to_go, the rewards from time t on."""
    S, A, R = traj[:3]
    togo = np.cumsum(R[::-1])[::-1]                  # togo[t] = r_t + r_{t+1} + ... + r_{T-1}
    g = np.zeros_like(theta)
    for t, (s, a) in enumerate(zip(S, A)):
        if a is None:                                 # the exit action is forced: no choice, no gradient
            continue
        g += grad_log_pi(theta, s, a) * (togo[t] if reward_to_go else sum(R))
    return g


def train(reward_to_go=False, seed=0, iters=200, n=8, lr=0.2):
    """REINFORCE: n trajectories per update. Returns the average return of each batch and the final logits."""
    rng = np.random.default_rng(seed)
    theta = np.zeros((len(STATES), 4))
    curve = []
    for _ in range(iters):
        trajs = [rollout(theta, rng) for _ in range(n)]
        curve.append(np.mean([sum(t[2]) for t in trajs]))
        theta += lr * np.mean([grad_estimate(theta, t, reward_to_go) for t in trajs], axis=0)
    return np.array(curve), theta


t0 = time.time()
curve, theta_trained = train(seed=1)
print(f'{time.time() - t0:.1f} s;  average return: first 10 updates {curve[:10].mean():.2f}, last 20 updates {curve[-20:].mean():.2f}')
go.Figure(go.Scatter(y=curve, line=dict(color=NAVY))).update_layout(xaxis_title='update', yaxis_title='average return of the batch',
                                                                     width=700, height=330, title='REINFORCE in the grid world')

0.6 s;  average return: first 10 updates -2.31, last 20 updates 0.50

policy_pair(theta0, theta_trained, ['uniform: every action 0.25', 'after REINFORCE (200 updates)'])

rng = np.random.default_rng(5)
trajs = [rollout(theta_trained, rng) for _ in range(8)]
t = torch.tensor(theta_trained, requires_grad=True)
loss = 0.0
for S, A, R, *_ in trajs:
    for s, a in zip(S, A):
        if a is not None:
            loss = loss - sum(R) * torch.log_softmax(t[IDX[s]], 0)[a] / len(trajs)
loss.backward()
by_hand = np.mean([grad_estimate(theta_trained, tr) for tr in trajs], axis=0)
print('max |(-autograd of the pseudo-loss) - hand-written estimate| =', np.abs(-t.grad.numpy() - by_hand).max())

max |(-autograd of the pseudo-loss) - hand-written estimate| = 1.3877787807814457e-16

rng = np.random.default_rng(7)
_, theta_mid = train(reward_to_go=True, seed=1, iters=15)            # a partly trained policy
for _ in range(1000):
    S, A, R, *_ = rollout(theta_mid, rng)
    if S[-1] == GOAL and set(S) & set(MUD) and len(S) <= 8:
        break
togo = np.cumsum(R[::-1])[::-1]
print('t      :', list(range(len(R))))
print('action :', [ACTIONS[a] if a is not None else 'exit' for a in A])
print('r_t    :', [round(r, 2) for r in R])
print('R(tau) :', [round(sum(R), 2)] * len(R), ' <- every action gets the same weight')
print('Q_hat_t:', [round(float(q), 2) for q in togo], ' <- reward-to-go: only the rewards from t on')

t      : [0, 1, 2, 3, 4, 5, 6]
action : ['up', 'right', 'right', 'left', 'up', 'right', 'exit']
r_t    : [-0.04, -0.04, -0.3, -0.3, -0.04, -0.04, 1.0]
R(tau) : [0.24, 0.24, 0.24, 0.24, 0.24, 0.24, 0.24]  <- every action gets the same weight
Q_hat_t: [0.24, 0.28, 0.32, 0.62, 0.92, 0.96, 1.0]  <- reward-to-go: only the rewards from t on

def compare_estimators(theta, n=4000, seed=0):
    rng = np.random.default_rng(seed)
    trajs = [rollout(theta, rng) for _ in range(n)]
    out = {}
    for rtg in (False, True):
        G = np.array([grad_estimate(theta, tr, rtg) for tr in trajs])
        out[rtg] = G
        name = 'reward-to-go Q_hat_t ' if rtg else 'whole return R(tau)  '
        print(f'{name} variance (summed over the 60 logits) {G.var(0).sum():7.1f}   mean gradient for (start, up) {G.mean(0)[IDX[START], 0]:+.3f}')
    return out

G = compare_estimators(theta0)

whole return R(tau)   variance (summed over the 60 logits)   215.4   mean gradient for (start, up) +0.082

reward-to-go Q_hat_t  variance (summed over the 60 logits)    91.4   mean gradient for (start, up) +0.078

t0 = time.time()
runs = {rtg: np.array([train(rtg, seed)[0] for seed in range(20)]) for rtg in (False, True)}
print(f'{time.time() - t0:.0f} s')
for rtg in (False, True):
    print('reward-to-go' if rtg else 'whole return', f': average return, updates 41-60 {runs[rtg][:, 40:60].mean():+.2f}, last 20 updates {runs[rtg][:, -20:].mean():+.2f}')
fig = go.Figure()
for rtg, col, name in ((False, GRAY, 'REINFORCE (whole return)'), (True, NAVY, 'REINFORCE with reward-to-go')):
    c = np.array([np.convolve(x, np.ones(10) / 10, mode='valid') for x in runs[rtg]])
    m, se = c.mean(0), c.std(0) / np.sqrt(len(c))
    x = np.arange(c.shape[1]) + 10
    fig.add_trace(go.Scatter(x=np.r_[x, x[::-1]], y=np.r_[m + 2 * se, (m - 2 * se)[::-1]], fill='toself', fillcolor=col, opacity=0.2,
                             line=dict(width=0), showlegend=False, hoverinfo='skip'))
    fig.add_trace(go.Scatter(x=x, y=m, name=name, line=dict(color=col, width=3)))
fig.update_layout(xaxis_title='update', yaxis_title='average return', width=760, height=380,
                  title='20 seeds, 8 trajectories per update (band: 2 standard errors)')
fig

20 s
whole return : average return, updates 41-60 -0.10, last 20 updates +0.42
reward-to-go : average return, updates 41-60 +0.03, last 20 updates +0.46

curve0, theta_stuck = train(seed=0)
print(f'seed 0: average return over the last 20 updates {curve0[-20:].mean():+.2f}')
print('pi(. | start):', {a: float(p) for a, p in zip(ACTIONS, policy(theta_stuck, START).round(2))})
policy_figure(theta_stuck, 'seed 0: stuck at the start cell')

seed 0: average return over the last 20 updates +0.09
pi(. | start): {'up': 0.04, 'right': 0.01, 'down': 0.02, 'left': 0.93}

for name, logits in (('uniform', [0, 0, 0, 0]), ('mostly up', [3, 0, 0, 0]), ('mostly right', [0, 3, 0, 0]), ('mostly left (stay put)', [0, 0, 0, 3])):
    th = np.zeros((len(STATES), 4))
    th[IDX[START]] = logits                       # 3 vs 0 gives probability 0.87; uniform everywhere else
    J, se = estimate_J(th, n=4000)
    print(f'start cell {name:23s} J ≈ {J:+.2f} ± {2 * se:.2f}')

start cell uniform                 J ≈ -2.64 ± 0.04

start cell mostly up               J ≈ -2.54 ± 0.04

start cell mostly right            J ≈ -2.78 ± 0.04

start cell mostly left (stay put)  J ≈ -2.37 ± 0.04

def train_baseline(baseline=None, seed=0, iters=200, n=8, lr=0.2):
    """REINFORCE with the whole return minus a baseline: None (plain), 'mean' (R - mean), or 'grpo' ((R - mean) / std)."""
    rng = np.random.default_rng(seed)
    theta = np.zeros((len(STATES), 4))
    curve = []
    for _ in range(iters):
        trajs = [rollout(theta, rng) for _ in range(n)]
        Rs = np.array([sum(t[2]) for t in trajs])
        curve.append(Rs.mean())
        w = Rs if baseline is None else Rs - Rs.mean()
        if baseline == 'grpo':
            w = w / (Rs.std() + 1e-8)
        grads = []
        for (S, A, *_), w_i in zip(trajs, w):
            g = np.zeros_like(theta)
            for s, a in zip(S, A):
                if a is not None:
                    g += grad_log_pi(theta, s, a) * w_i
            grads.append(g)
        theta += lr * np.mean(grads, axis=0)          # baseline=None is exactly train() from Demo 1
    return np.array(curve), theta


fig = go.Figure()
for baseline, col, name in ((None, GRAY, 'plain: R'), ('mean', NAVY, 'baseline: R − mean'), ('grpo', GOLD, 'GRPO-style: (R − mean) / std')):
    for seed, dash in ((0, 'solid'), (6, 'dot')):
        c, th = train_baseline(baseline, seed)
        print(f'{name:30s} seed {seed}: last 20 updates {c[-20:].mean():+.2f}')
        fig.add_trace(go.Scatter(y=np.convolve(c, np.ones(10) / 10, mode='valid'), name=f'{name}, seed {seed}',
                                 line=dict(color=col, width=3, dash=dash)))
fig.update_layout(xaxis_title='update', yaxis_title='average return (10-update average)', width=760, height=400,
                  title='The two stuck seeds, with and without a baseline')

plain: R                       seed 0: last 20 updates +0.09
plain: R                       seed 6: last 20 updates -0.21

baseline: R − mean             seed 0: last 20 updates +0.43
baseline: R − mean             seed 6: last 20 updates +0.46
GRPO-style: (R − mean) / std   seed 0: last 20 updates +0.54

GRPO-style: (R − mean) / std   seed 6: last 20 updates +0.46

STEP_COST, MUD_COST = 0.0, 0.0            # experiment 1: reward only at the end
G = compare_estimators(theta0, n=1000)
print('estimates identical:', np.allclose(G[False], G[True]))
STEP_COST, MUD_COST = -0.04, -0.3          # restore the lecture's world

whole return R(tau)   variance (summed over the 60 logits)     9.3   mean gradient for (start, up) +0.046
reward-to-go Q_hat_t  variance (summed over the 60 logits)     9.3   mean gradient for (start, up) +0.046
estimates identical: True

Notebook section	Slides
A Slippery Grid World	A Grid World Example, Random Transitions, Markov Decision Processes
Policies	Policies, Neural Networks as Policies
Trajectories and Their Probability	Trajectories, The Probability of a Trajectory, The Objective
REINFORCE (Lecture Demo 1)	The Log-Derivative Trick, The Environment Drops Out, REINFORCE, REINFORCE Is Weighted Fine-Tuning
Reward-to-Go (Lecture Demo 2)	Credit Assignment, Causality, REINFORCE with Reward-to-Go, Reward-to-Go in the Grid World
When REINFORCE Gets Stuck	(a preview of Lecture 22: baselines)
Try It Yourself

Lecture 21: Reinforcement Learning in a Grid World (CS 189, Fall 2026)

A Slippery Grid World

Policies

Trajectories and Their Probability

(Lecture Demo 1) REINFORCE

(Lecture Demo 2) Reward-to-Go

When REINFORCE Gets Stuck

Try It Yourself