"""Chapter 3 lab: Markov decision processes as numpy arrays. Your tasks, in order (each one has tests in test_mdp.py): Task 1 build_gridworld fill in P[s, a, s'] and R[s, a] for the slippery gridworld Task 2 policy_matrices collapse the MDP under a policy: P_pi[s, s'] and r_pi[s] Task 3 evaluate_exact V^pi = (I - gamma P_pi)^{-1} r_pi (one linear solve) Task 4 evaluate_iterative V_{k+1} = r_pi + gamma P_pi V_k (repeat until it stops changing) Task 5 q_from_v Q^pi(s, a) = R[s, a] + gamma sum_s' P[s, a, s'] V^pi(s') Conventions (same as the chapter and its demo; see envs.py): * P has shape (S, A, S): P[s, a, s2] = probability of landing in s2. * R has shape (S, A): R[s, a] = expected reward for taking a in s. * pi has shape (S, A): pi[s, a] = probability that the policy picks a in s. * Gridworld states are numbered s = row * N_COLS + col; actions are 0 up, 1 right, 2 down, 3 left. Run the tests with: python -m pytest -q """ from __future__ import annotations import numpy as np from envs import (EXIT_REWARD, N_ACTIONS, N_STATES, is_terminal, kind, move, perpendicular) # --------------------------------------------------------------------------- # Task 1 # --------------------------------------------------------------------------- def build_gridworld(slip: float = 0.2, step_reward: float = -0.04) -> tuple[np.ndarray, np.ndarray]: """Build the gridworld MDP of the chapter as arrays. Dynamics. For every non-terminal state s and action a: * with probability 1 - slip you go where you intended: move(s, a) * with probability slip / 2 each you slip sideways: move(s, b) for each b in perpendicular(a) move() already handles bumping into the wall or the edge (you stay put). Two of these outcomes can land in the SAME cell (e.g. two bumps), so accumulate probabilities with += rather than overwrite them. Rewards. Every action costs `step_reward` (a negative number = a cost). On top of that, the step that ENTERS the goal earns +1 and the step that enters the pit earns -1 (EXIT_REWARD maps 'G' -> 1.0, 'P' -> -1.0; use EXIT_REWARD.get(kind(s2), 0.0)). R[s, a] is the EXPECTED reward: R[s, a] = step_reward + sum_{s2} P[s, a, s2] * EXIT_REWARD.get(kind(s2), 0.0) Terminal states (is_terminal(s) is True for goal, pit and the dummy wall cell) are absorbing: every action leads back to s with probability 1 and reward 0. That makes the episodic task a continuing one whose rewards are all zero after the episode ends (see the chapter). Returns: P: array of shape (N_STATES, N_ACTIONS, N_STATES), every P[s, a] sums to 1 R: array of shape (N_STATES, N_ACTIONS) """ P = np.zeros((N_STATES, N_ACTIONS, N_STATES)) R = np.zeros((N_STATES, N_ACTIONS)) for s in range(N_STATES): for a in range(N_ACTIONS): # TODO: fill in P[s, a, :] and R[s, a] as described above. raise NotImplementedError("Task 1: build_gridworld") return P, R # --------------------------------------------------------------------------- # Task 2 # --------------------------------------------------------------------------- def policy_matrices(P: np.ndarray, R: np.ndarray, pi: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """Average the MDP over the policy's action choice. P_pi[s, s2] = sum_a pi[s, a] * P[s, a, s2] (a Markov chain over states) r_pi[s] = sum_a pi[s, a] * R[s, a] (expected reward per step) Args: P: (S, A, S) transition array. R: (S, A) reward array. pi: (S, A) policy. Returns: P_pi with shape (S, S) (each row sums to 1) and r_pi with shape (S,). Hint: np.einsum("sa,sat->st", pi, P) does the first sum in one line; a double loop is also fine. """ # TODO raise NotImplementedError("Task 2: policy_matrices") # --------------------------------------------------------------------------- # Task 3 # --------------------------------------------------------------------------- def evaluate_exact(P: np.ndarray, R: np.ndarray, pi: np.ndarray, gamma: float) -> np.ndarray: """Policy evaluation by one linear solve (the matrix Bellman equation). V = r_pi + gamma P_pi V => (I - gamma P_pi) V = r_pi Use policy_matrices, then np.linalg.solve (not np.linalg.inv: solving is faster and more accurate than forming the inverse). Requires gamma < 1. Returns: V with shape (S,): V[s] = V^pi(s). """ # TODO raise NotImplementedError("Task 3: evaluate_exact") # --------------------------------------------------------------------------- # Task 4 # --------------------------------------------------------------------------- def evaluate_iterative(P: np.ndarray, R: np.ndarray, pi: np.ndarray, gamma: float, tol: float = 1e-10, max_iters: int = 100_000) -> tuple[np.ndarray, int]: """Policy evaluation by iterating the Bellman expectation equation. Start from V_0 = 0 (all zeros) and repeat V_{k+1} = r_pi + gamma P_pi V_k until max_s |V_{k+1}(s) - V_k(s)| < tol. Return the last V_{k+1} and the number of updates k+1 you performed (so the first update counts as 1). If tol is never reached, stop after max_iters updates. Returns: (V, n_iters): V with shape (S,), n_iters an int. """ # TODO raise NotImplementedError("Task 4: evaluate_iterative") # --------------------------------------------------------------------------- # Task 5 # --------------------------------------------------------------------------- def q_from_v(P: np.ndarray, R: np.ndarray, V: np.ndarray, gamma: float) -> np.ndarray: """Action values from state values (one step of lookahead): Q[s, a] = R[s, a] + gamma * sum_{s2} P[s, a, s2] * V[s2] Args: P: (S, A, S). R: (S, A). V: (S,) state values of some policy pi. Returns: Q with shape (S, A). If V = V^pi then Q = Q^pi, and sum_a pi[s, a] Q[s, a] = V[s] for every s (the tests check this). Hint: P @ V has shape (S, A). """ # TODO raise NotImplementedError("Task 5: q_from_v")