"""Chapter 4 lab: dynamic programming. YOUR FILE: fill in the five TODOs, in order. All functions work on an MDP given as arrays (see envs.py): P: (S, A, S) transition probabilities, R: (S, A) expected rewards. A policy is an (S, A) array whose rows are probability distributions (a deterministic policy is one-hot: exactly one 1.0 per row). Run the tests with: python -m pytest -q Each task has its own tests (test_task1_..., test_task2_..., and so on), so you can work through them one at a time: python -m pytest -q -k task1 """ import numpy as np # --------------------------------------------------------------------------- # Task 1: one-step lookahead # --------------------------------------------------------------------------- def q_from_v(P, R, V, gamma): """One-step lookahead: the action values implied by a state-value vector V. Q[s, a] = R[s, a] + gamma * sum_{s2} P[s, a, s2] * V[s2] This is the "look one step ahead, then trust V" computation inside every Bellman backup in the chapter. Every other task calls it. Args: P: (S, A, S) transition probabilities. R: (S, A) expected rewards. V: (S,) state values. gamma: discount factor in [0, 1]. Returns: Q: (S, A) array. Hint: numpy's `P @ V` contracts the last axis of P with V, giving an (S, A) array. One line is enough. """ # TODO: return the (S, A) array of one-step lookahead values. raise NotImplementedError("Task 1: q_from_v") # --------------------------------------------------------------------------- # Task 2: iterative policy evaluation # --------------------------------------------------------------------------- def policy_evaluation(P, R, policy, gamma, tol=1e-10, max_sweeps=100_000, V0=None): """Iterative policy evaluation (synchronous sweeps of the Bellman expectation backup). V_{k+1}(s) = sum_a pi(a|s) * ( R[s, a] + gamma * sum_{s2} P[s, a, s2] V_k(s2) ) Start from V0 (zeros if None). Stop after the first sweep whose largest change max_s |V_{k+1}(s) - V_k(s)| is below `tol` (or after `max_sweeps`). "Synchronous" means every V_{k+1}(s) is computed from the OLD vector V_k (compute the whole new vector, then replace the old one). Returns: V: (S,) the final value estimate V_{k+1}. n_sweeps: number of sweeps performed (including the last one). Hint: the bracket is q_from_v(P, R, V, gamma); averaging it over pi(a|s) is an elementwise product with `policy` followed by a sum over axis 1. """ S = P.shape[0] V = np.zeros(S) if V0 is None else np.array(V0, dtype=float) # TODO: loop over sweeps; compute V_new from V; track the largest change; # stop when it is below tol. Return (V, number_of_sweeps). raise NotImplementedError("Task 2: policy_evaluation") # --------------------------------------------------------------------------- # Task 3: greedy policy # --------------------------------------------------------------------------- def greedy_policy(P, R, V, gamma): """The deterministic policy that is greedy with respect to V. pi(s) = argmax_a R[s, a] + gamma * sum_{s2} P[s, a, s2] V[s2] Ties go to the lowest action index (what np.argmax does). Returns: policy: (S, A) one-hot array (exactly one 1.0 per row). """ # TODO: compute Q with q_from_v, take the argmax per row, build a one-hot array. raise NotImplementedError("Task 3: greedy_policy") # --------------------------------------------------------------------------- # Task 4: value iteration # --------------------------------------------------------------------------- def value_iteration(P, R, gamma, tol=1e-10, max_sweeps=100_000, V0=None): """Value iteration: repeat the Bellman optimality backup until values stop changing. V_{k+1}(s) = max_a R[s, a] + gamma * sum_{s2} P[s, a, s2] V_k(s2) Start from V0 (zeros if None). After each sweep record delta_k = max_s |V_{k+1}(s) - V_k(s)|. Stop after the first sweep with delta_k < tol (or after max_sweeps). Returns: V: (S,) final values. policy: (S, A) one-hot greedy policy with respect to the final V. deltas: list of floats, one delta per sweep (len = number of sweeps). """ S = P.shape[0] V = np.zeros(S) if V0 is None else np.array(V0, dtype=float) deltas = [] # TODO: the loop. It is policy_evaluation with "average over pi" replaced # by "max over a". Append each sweep's delta (as a Python float). raise NotImplementedError("Task 4: value_iteration") # --------------------------------------------------------------------------- # Task 5: policy iteration # --------------------------------------------------------------------------- def policy_iteration(P, R, gamma, policy0=None, tol=1e-10, max_iters=1_000): """Policy iteration: evaluate, improve greedily, repeat until the policy is stable. 1. Evaluate the current policy with policy_evaluation(...) (use `tol`). 2. Compute Q = q_from_v(P, R, V, gamma). In each state, switch to the best action, but ONLY if it beats the current action by more than 1e-9 (otherwise round-off can make the policy flip between equally good actions forever). 3. If no state switched, stop. Args: policy0: (S, A) one-hot starting policy; None means "always action 0". Returns: policy: (S, A) one-hot final policy. V: (S,) its value V^pi. history: list of the policies evaluated, history[0] = policy0 and history[-1] = the final policy (len = number of evaluations). """ S, A = R.shape if policy0 is None: policy0 = np.zeros((S, A)) policy0[:, 0] = 1.0 policy = np.array(policy0, dtype=float) history = [] for _ in range(max_iters): history.append(policy.copy()) # TODO: (1) evaluate `policy`; (2) compute Q and, per state, decide whether # to switch from the current action np.argmax(policy, axis=1) to # np.argmax(Q, axis=1) (only if better by > 1e-9); # (3) if nothing switched, break; otherwise build the new one-hot policy. raise NotImplementedError("Task 5: policy_iteration") return policy, V, history