"""Environments and helpers for the Chapter 5 lab. This file is complete: you do not need to edit it. Read it once so you know what `env.reset()` and `env.step(a)` return. Two environments: * `RandomWalk`: the 5-state random walk used for prediction (MC vs. TD). * `CliffWalking`: the 4 x 12 cliff used for control (SARSA vs. Q-learning). And a few helpers that the tests and `run.py` use: `generate_episode`, `rms_error`, `prediction_curve`, `train`, `greedy_path`, `render_path`. """ from __future__ import annotations import numpy as np # --------------------------------------------------------------------------- # Prediction: the 5-state random walk # --------------------------------------------------------------------------- class RandomWalk: """The 5-state random walk (Sutton & Barto, Example 6.2). [exit, reward 0] A B C D E [exit, reward 1] 0 1 2 3 4 (state indices) Every episode starts in C (state 2). Each step moves left or right with probability 1/2. Leaving E to the right gives reward +1; every other step gives 0. There are no actions: the "policy" is the coin flip, so this is a pure prediction problem. With gamma = 1, the true value of each state is the probability of exiting on the right: V(A..E) = 1/6, 2/6, 3/6, 4/6, 5/6 (`RandomWalk.true_values`). """ n_states = 5 names = "ABCDE" true_values = np.arange(1, 6) / 6.0 def __init__(self, seed: int = 0): self.rng = np.random.default_rng(seed) self.s = 2 def reset(self) -> int: self.s = 2 return self.s def step(self): """Move one step. Returns (s_next, r, done); s_next is None when done.""" s2 = self.s + (1 if self.rng.random() < 0.5 else -1) if s2 < 0: return None, 0.0, True if s2 >= self.n_states: return None, 1.0, True self.s = s2 return s2, 0.0, False def generate_episode(env: RandomWalk) -> list[tuple]: """Run one episode and return its transitions. Returns a list of tuples (s_t, r_t, s_next, done), one per step, where r_t is the reward received on leaving s_t (the course's convention) and s_next is None on the final, terminating step. Example: C -> D -> E -> exit right gives [(2, 0.0, 3, False), (3, 0.0, 4, False), (4, 1.0, None, True)] """ episode = [] s = env.reset() while True: s_next, r, done = env.step() episode.append((s, r, s_next, done)) if done: return episode s = s_next def rms_error(V: np.ndarray, true_values: np.ndarray) -> float: """Root-mean-square error between an estimate and the true values.""" return float(np.sqrt(np.mean((np.asarray(V) - true_values) ** 2))) def prediction_curve(update_fn, alpha: float, n_episodes: int = 100, n_runs: int = 100, seed: int = 0) -> np.ndarray: """Average RMS error of a prediction method over many independent runs. `update_fn(V, episode, alpha, gamma)` is your `mc_update` or `td0_update`. Each run starts from V = 0.5 everywhere and processes `n_episodes` episodes. Returns an array of length n_episodes + 1: entry k is the RMS error after k episodes, averaged over runs. """ errors = np.zeros(n_episodes + 1) for run in range(n_runs): env = RandomWalk(seed=seed + run) V = np.full(RandomWalk.n_states, 0.5) errors[0] += rms_error(V, RandomWalk.true_values) for k in range(1, n_episodes + 1): update_fn(V, generate_episode(env), alpha, 1.0) errors[k] += rms_error(V, RandomWalk.true_values) return errors / n_runs # --------------------------------------------------------------------------- # Control: cliff walking # --------------------------------------------------------------------------- class CliffWalking: """Cliff walking (Sutton & Barto, Example 6.6). row 0 . . . . . . . . . . . . row 1 . . . . . . . . . . . . row 2 . . . . . . . . . . . . row 3 S C C C C C C C C C C G S = start, G = goal, C = cliff. States are numbered row * 12 + col, so S = 36 and G = 47. Actions: 0 = up, 1 = right, 2 = down, 3 = left. Moving into a wall leaves you where you are. Rewards: -1 per step. Stepping into the cliff gives -100 and sends you back to S; the episode continues. Reaching G ends the episode. The shortest path (along the cliff edge) takes 13 steps: return -13. """ rows, cols = 4, 12 n_states = rows * cols n_actions = 4 start, goal = 36, 47 moves = [(-1, 0), (0, 1), (1, 0), (0, -1)] def __init__(self): self.s = self.start @classmethod def is_cliff(cls, s: int) -> bool: r, c = divmod(s, cls.cols) return r == cls.rows - 1 and 0 < c < cls.cols - 1 def reset(self) -> int: self.s = self.start return self.s def step(self, a: int): """Take action a. Returns (s_next, r, done).""" r, c = divmod(self.s, self.cols) dr, dc = self.moves[a] r = min(self.rows - 1, max(0, r + dr)) c = min(self.cols - 1, max(0, c + dc)) s2 = r * self.cols + c if self.is_cliff(s2): self.s = self.start return self.s, -100.0, False self.s = s2 return s2, -1.0, s2 == self.goal def train(episode_fn, env, n_episodes: int, alpha: float, epsilon: float, gamma: float = 1.0, seed: int = 0): """Train a control agent from Q = 0. `episode_fn(env, Q, alpha, epsilon, gamma, rng)` is your `sarsa_episode` or `q_learning_episode`: it runs one episode, updates Q in place and returns the episode's total reward. Returns (Q, returns) with Q of shape (n_states, n_actions) and returns of shape (n_episodes,). """ rng = np.random.default_rng(seed) Q = np.zeros((env.n_states, env.n_actions)) returns = np.empty(n_episodes) for k in range(n_episodes): returns[k] = episode_fn(env, Q, alpha, epsilon, gamma, rng) return Q, returns def greedy_path(Q: np.ndarray, env: CliffWalking | None = None, max_len: int = 100): """Follow argmax_a Q(s, a) from the start, with no exploration. Returns (states, reached_goal). Stops at the goal, on falling into the cliff, on revisiting a state (the greedy policy loops), or after max_len. """ env = env or CliffWalking() s = env.reset() states, seen = [s], {s} for _ in range(max_len): s2, r, done = env.step(int(np.argmax(Q[s]))) if r <= -100: return states, False states.append(s2) if done: return states, True if s2 in seen: return states, False seen.add(s2) s = s2 return states, False def render_path(Q: np.ndarray) -> str: """ASCII picture of the greedy path: '*' on the path, 'C' cliff, S/G start/goal.""" env = CliffWalking() states, _ = greedy_path(Q, env) on_path = set(states) lines = [] for r in range(env.rows): row = [] for c in range(env.cols): s = r * env.cols + c if s == env.start: row.append("S") elif s == env.goal: row.append("G") elif env.is_cliff(s): row.append("C") elif s in on_path: row.append("*") else: row.append(".") lines.append(" ".join(row)) return "\n".join(lines)