Monte Carlo Control

Complete ε-greedy MC control on the same GridWorld (the step function is unchanged from the prediction section):

import numpy as np
from collections import defaultdict

ROWS, COLS = 4, 4
WALL, GOAL, START = (1, 1), (3, 3), (0, 0)
ACTIONS = ['up', 'down', 'left', 'right']
MOVES = {'up': (-1, 0), 'down': (1, 0), 'left': (0, -1), 'right': (0, 1)}
GAMMA = 0.9


def step(state, action):
    """One environment step. Returns (next_state, reward, done)."""
    dr, dc = MOVES[action]
    nxt = (state[0] + dr, state[1] + dc)
    if not (0 <= nxt[0] < ROWS and 0 <= nxt[1] < COLS) or nxt == WALL:
        nxt = state  # bounce off wall or edge, still pay the step cost
    reward = -1.0 + (10.0 if nxt == GOAL else 0.0)
    return nxt, reward, nxt == GOAL


def epsilon_greedy(Q, state, epsilon, rng):
    """Pick a random action with prob. epsilon, else the greedy one."""
    if rng.random() < epsilon:
        return ACTIONS[rng.integers(4)]
    qs = [Q[(state, a)] for a in ACTIONS]
    best = max(qs)
    # break ties randomly so early learning doesn't lock onto one action
    candidates = [a for a, q in zip(ACTIONS, qs) if q == best]
    return candidates[rng.integers(len(candidates))]


def mc_control(num_episodes, epsilon=0.1, gamma=GAMMA, seed=0, max_steps=1000):
    """First-visit Monte Carlo control with an epsilon-greedy policy."""
    rng = np.random.default_rng(seed)
    Q = defaultdict(float)   # action-value estimates
    N = defaultdict(int)     # first-visit counts per (state, action)
    episode_returns = []
    for _ in range(num_episodes):
        # 1. Generate an episode with the CURRENT epsilon-greedy policy
        state, trajectory, done, t = START, [], False, 0
        while not done and t < max_steps:
            action = epsilon_greedy(Q, state, epsilon, rng)
            next_state, reward, done = step(state, action)
            trajectory.append((state, action, reward))
            state = next_state
            t += 1
        # 2. Walk backward, computing returns
        G = 0.0
        returns = []
        for state, action, reward in reversed(trajectory):
            G = reward + gamma * G
            returns.append((state, action, G))
        returns.reverse()
        episode_returns.append(returns[0][2])  # return from the start state
        # 3. First-visit update for each (state, action) pair
        seen = set()
        for state, action, G in returns:
            if (state, action) not in seen:
                seen.add((state, action))
                N[(state, action)] += 1
                Q[(state, action)] += (G - Q[(state, action)]) / N[(state, action)]
        # (No explicit improvement step: epsilon_greedy reads the updated Q
        #  on the very next episode. Evaluation and improvement interleave.)
    return Q, episode_returns


Q, episode_returns = mc_control(num_episodes=20_000)

ARROW = {'up': '^', 'down': 'v', 'left': '<', 'right': '>'}
print("Greedy policy after 20,000 episodes:")
for r in range(ROWS):
    row = []
    for c in range(COLS):
        if (r, c) == WALL:
            row.append('X')
        elif (r, c) == GOAL:
            row.append('G')
        else:
            row.append(ARROW[max(ACTIONS, key=lambda a: Q[((r, c), a)])])
    print('  ' + ' '.join(row))

print("\nV(s) = max_a Q(s,a):")
for r in range(ROWS):
    print('  '.join(
        '  WALL' if (r, c) == WALL else '  GOAL' if (r, c) == GOAL
        else f'{max(Q[((r, c), a)] for a in ACTIONS):6.2f}' for c in range(COLS)))

print(f"\nMean return, first 100 episodes: {np.mean(episode_returns[:100]):.2f}")
print(f"Mean return, last 100 episodes:  {np.mean(episode_returns[-100:]):.2f}")

Output:

Greedy policy after 20,000 episodes:
  > > v v
  v X > v
  > > > v
  > ^ > G

V(s) = max_a Q(s,a):
  0.62    1.93    3.37    5.01
  1.66    WALL    4.97    6.87
  3.17    4.85    6.90    9.00
  1.86    2.56    9.00    GOAL

Mean return, first 100 episodes: 0.00
Mean return, last 100 episodes:  0.69