"""Reinforcement learning on a 4 x 4 gridworld: the value function of a policy is the fixed point of its Bellman operator, the greedy policy is optimal, and Q-learning finds the same policy from interaction alone. Purpose ------- Numerical companion to the glossary entry 'reinforcementlearning' (reinforcement learning (RL)). The gridworld is the one of the entry 'gridworld': a 4 x 4 lattice of cells (row, col) with row 0 at the bottom, a wall at cell (2, 2), a goal at cell (3, 3), and the four moves up, down, left, right. A move that would leave the grid or enter the wall leaves the agent in place. Every move yields reward -1, except the move that enters the goal, which yields +1 and ends the episode (the goal is absorbing). No discounting (gamma = 1): every episode ends after finitely many moves under every policy considered here. The demo checks the entry's claims: (1) the state-value function of the equiprobable random policy is the fixed point of the policy's Bellman operator, reached by iterating it from zero (iterative policy evaluation); (2) the greedy policy with respect to that value function already reaches the goal on a shortest path from every cell, and its value function is the optimal one, 2 - d(s) with d(s) the shortest path length to the goal; (3) Q-learning with random moves for exploration, which never uses the transition function and learns only from the rewards of the actions it tried, arrives at the same optimal policy, and its return per episode rises from that of a random walk to the optimum -- delayed reward at work, since the +1 arrives only at the last move. Deterministic: fixed seed for the exploration noise. Self-contained: numpy + matplotlib only. Blocks ------ [B-gridworld] Build the MDP: 15 states (16 cells minus the wall), 4 actions, deterministic transition function, reward -1 per move and +1 for entering the goal; check the boundary and wall rules. [B-evaluate] Iterative policy evaluation of the random policy: apply the Bellman operator of the policy from v = 0 until the change is below 1e-10; check the result is a fixed point, that the iteration converged, and that the values are negative and decrease with the distance to the goal. [B-improve] Greedy policy with respect to the random policy's value function: check it reaches the goal from every cell in the shortest number of moves, and that its value function is the optimal one, 2 - d(s). [B-qlearning] Q-learning (exploration probability 0.1, update weight 0.5) over 400 episodes from random start cells, learning from the rewards of the actions tried only; check that the greedy policy of the learned Q equals the optimal policy and that the average return of the last 50 episodes exceeds that of the first 50. [B-plot] Preview: the gridworld with the optimal values and policy arrows, and the return per episode of Q-learning. Outputs ------- reinforcementlearning_values.csv : row, col, v_random, v_optimal, action (optimal move) per non-wall cell reinforcementlearning_returns.csv : episode, return, running mean over 20 episodes (Q-learning) reinforcementlearning.png : preview (checking only) """ import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from pathlib import Path OUT_DIR = Path(__file__).parent report = [] def check(name, ok): report.append((name, bool(ok))) print(f" [{'ok' if ok else 'FAIL'}] {name}") # ---- [B-gridworld] the MDP of the gridworld entry H, W = 4, 4 WALL, GOAL = (2, 2), (3, 3) ACTIONS = {"up": (1, 0), "down": (-1, 0), "left": (0, -1), "right": (0, 1)} STATES = [(r, c) for r in range(H) for c in range(W) if (r, c) != WALL] def step(s, a): """Next state and reward of a deterministic move; the goal is absorbing.""" if s == GOAL: return s, 0.0 dr, dc = ACTIONS[a] n = (s[0] + dr, s[1] + dc) if not (0 <= n[0] < H and 0 <= n[1] < W) or n == WALL: n = s return n, (1.0 if n == GOAL else -1.0) check("[B-gridworld] 15 states and 4 actions", len(STATES) == 15 and len(ACTIONS) == 4) check("[B-gridworld] a move into the boundary or the wall leaves the agent in place", step((0, 0), "down")[0] == (0, 0) and step((1, 2), "up")[0] == (1, 2)) check("[B-gridworld] entering the goal yields +1, any other move -1", step((3, 2), "right") == (GOAL, 1.0) and step((0, 0), "up") == ((1, 0), -1.0)) # ---- [B-evaluate] the value function of the random policy as a fixed point def bellman_policy(v, policy): """One application of the Bellman operator of a policy (gamma = 1): (T v)(s) = sum_a pi(a|s) [r(s,a) + v(s')].""" out = {} for s in STATES: if s == GOAL: out[s] = 0.0 continue out[s] = sum(p * (step(s, a)[1] + v[step(s, a)[0]]) for a, p in policy[s].items()) return out random_policy = {s: {a: 0.25 for a in ACTIONS} for s in STATES} v = {s: 0.0 for s in STATES} sweeps, deltas = 0, [] while True: v_new = bellman_policy(v, random_policy) delta = max(abs(v_new[s] - v[s]) for s in STATES) deltas.append(delta); v = v_new; sweeps += 1 if delta < 1e-10 or sweeps > 10_000: break v_random = v print(f" policy evaluation: {sweeps} sweeps; v_random(0,0) = {v_random[(0, 0)]:.2f}, " f"v_random(3,2) = {v_random[(3, 2)]:.2f}") check("[B-evaluate] the iteration converged (change below 1e-10)", deltas[-1] < 1e-10) check("[B-evaluate] the result is a fixed point of the policy's Bellman operator", max(abs(bellman_policy(v_random, random_policy)[s] - v_random[s]) for s in STATES) < 1e-8) def shortest(s): """Shortest number of moves from s to the goal (breadth-first search).""" from collections import deque seen, q = {s: 0}, deque([s]) while q: u = q.popleft() if u == GOAL: return seen[u] for a in ACTIONS: n = step(u, a)[0] if n not in seen: seen[n] = seen[u] + 1; q.append(n) return None d = {s: shortest(s) for s in STATES} check("[B-evaluate] the random policy's values are negative away from the goal " "and decrease with the distance to it", all(v_random[s] < 0 for s in STATES if s != GOAL) and all(v_random[s] >= v_random[t] for s in STATES for t in STATES if d[s] < d[t])) # ---- [B-improve] the greedy policy is optimal def greedy(vfun): pol = {} for s in STATES: if s == GOAL: pol[s] = None; continue pol[s] = max(ACTIONS, key=lambda a: step(s, a)[1] + vfun[step(s, a)[0]]) return pol greedy_policy = greedy(v_random) def rollout(pol, s, limit=100): """Return and number of moves when following pol from s.""" G, n = 0.0, 0 while s != GOAL and n < limit: s, r = step(s, pol[s]); G += r; n += 1 return G, n check("[B-improve] the greedy policy reaches the goal on a shortest path from every cell", all(rollout(greedy_policy, s)[1] == d[s] for s in STATES if s != GOAL)) v_greedy = {s: rollout(greedy_policy, s)[0] for s in STATES} v_optimal = {s: (0.0 if s == GOAL else 2.0 - d[s]) for s in STATES} check("[B-improve] its value function is the optimal one, 2 - d(s)", all(abs(v_greedy[s] - v_optimal[s]) < 1e-12 for s in STATES)) check("[B-improve] the optimal values are a fixed point of the optimal Bellman " "operator", all(abs(max(step(s, a)[1] + v_optimal[step(s, a)[0]] for a in ACTIONS) - v_optimal[s]) < 1e-12 for s in STATES if s != GOAL)) # ---- [B-qlearning] learning from interaction only rng = np.random.default_rng(0) EPS, ALPHA_Q, EPISODES = 0.1, 0.5, 400 Q = {s: {a: 0.0 for a in ACTIONS} for s in STATES} starts = [s for s in STATES if s != GOAL] returns = [] for ep in range(EPISODES): s = starts[rng.integers(len(starts))] G, n = 0.0, 0 while s != GOAL and n < 200: if rng.random() < EPS: a = list(ACTIONS)[rng.integers(4)] else: a = max(ACTIONS, key=lambda x: Q[s][x]) n_s, r = step(s, a) target = r + (0.0 if n_s == GOAL else max(Q[n_s].values())) Q[s][a] += ALPHA_Q * (target - Q[s][a]) G += r; s = n_s; n += 1 returns.append(G) returns = np.array(returns) q_policy = {s: (None if s == GOAL else max(ACTIONS, key=lambda a: Q[s][a])) for s in STATES} check("[B-qlearning] the greedy policy of the learned Q reaches the goal on a " "shortest path from every cell", all(rollout(q_policy, s)[1] == d[s] for s in STATES if s != GOAL)) first, last = returns[:50].mean(), returns[-50:].mean() print(f" Q-learning: mean return first 50 episodes {first:.1f}, last 50 episodes {last:.1f}" f" (optimum averaged over start cells {np.mean([v_optimal[s] for s in starts]):.1f})") check("[B-qlearning] the average return of the last 50 episodes exceeds that of " "the first 50", last > first) # ---- CSVs with open(OUT_DIR / "reinforcementlearning_values.csv", "w") as f: f.write("row,col,v_random,v_optimal,action\n") for (r, c) in STATES: f.write(f"{r},{c},{v_random[(r, c)]:.2f},{v_optimal[(r, c)]:.0f}," f"{greedy_policy[(r, c)] or 'goal'}\n") running = np.convolve(returns, np.ones(20) / 20, mode="full")[:EPISODES] running[:19] = np.nan with open(OUT_DIR / "reinforcementlearning_returns.csv", "w") as f: f.write("episode,ret,mean20\n") for i, (g, m) in enumerate(zip(returns, running), 1): f.write(f"{i},{g:.0f},{'nan' if np.isnan(m) else f'{m:.2f}'}\n") # ---- [B-plot] preview fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 4.2)) ARROW = {"up": (0, 0.3), "down": (0, -0.3), "left": (-0.3, 0), "right": (0.3, 0)} for r in range(H): for c in range(W): if (r, c) == WALL: ax1.add_patch(plt.Rectangle((c, r), 1, 1, color="0.6")) continue if (r, c) == GOAL: ax1.add_patch(plt.Rectangle((c, r), 1, 1, color="#c8e6c9")) ax1.text(c + 0.5, r + 0.5, "goal", ha="center", va="center", fontsize=9) continue ax1.text(c + 0.5, r + 0.78, f"{v_random[(r, c)]:.1f}", ha="center", va="center", fontsize=7, color="0.35") ax1.text(c + 0.5, r + 0.22, f"{v_optimal[(r, c)]:.0f}", ha="center", va="center", fontsize=9) dx, dy = ARROW[greedy_policy[(r, c)]] ax1.annotate("", xy=(c + 0.5 + dx, r + 0.5 + dy), xytext=(c + 0.5, r + 0.5), arrowprops=dict(arrowstyle="->", lw=1.2)) ax1.set_xlim(0, W); ax1.set_ylim(0, H); ax1.set_xticks(range(W + 1)); ax1.set_yticks(range(H + 1)) ax1.grid(True, color="black", lw=0.6); ax1.set_aspect("equal") ax1.set_xlabel("column"); ax1.set_ylabel("row") ax1.set_title("Values of the random policy (small) and optimal values with policy") ax2.plot(np.arange(1, EPISODES + 1), returns, ".", ms=3, color="0.6", label="return of the episode") ax2.plot(np.arange(1, EPISODES + 1), running, "k-", lw=1.5, label="mean over 20 episodes") ax2.axhline(np.mean([v_optimal[s] for s in starts]), ls="--", color="black", lw=1, label="optimal (mean over start cells)") ax2.set_xlabel("episode"); ax2.set_ylabel("return") ax2.set_title("Q-learning: return per episode") ax2.legend(frameon=False, fontsize=8) fig.tight_layout() fig.savefig(OUT_DIR / "reinforcementlearning.png", dpi=110) n_ok = sum(ok for _, ok in report) print(f"\n{n_ok}/{len(report)} checks pass") if n_ok != len(report): raise SystemExit(1)