"""Chapter 5 lab: Monte Carlo and TD learning. Fill in the five tasks below, in order. Each one is a few lines. Run the tests for one task at a time, for example: python -m pytest -q -k mc_update # Task 1 python -m pytest -q -k td0 # Task 2 python -m pytest -q -k epsilon_greedy # Task 3 python -m pytest -q -k target # Task 4 python -m pytest -q -k control # Task 5 Conventions (the same as the chapter): * r_t is the reward received on leaving s_t (after taking a_t). * A prediction episode is a list of (s_t, r_t, s_next, done) tuples; s_next is None on the last step. See envs.generate_episode. * Q is a float array of shape (n_states, n_actions); V has shape (n_states,). * Every update function modifies V or Q IN PLACE and returns nothing. """ from __future__ import annotations import numpy as np # --------------------------------------------------------------------------- # Part A: prediction (learning V for a fixed policy) # --------------------------------------------------------------------------- def mc_update(V: np.ndarray, episode: list[tuple], alpha: float, gamma: float = 1.0) -> None: """Task 1. Every-visit, constant-alpha Monte Carlo update, applied after the episode. Inputs V float array of shape (n_states,), modified IN PLACE. episode list of (s_t, r_t, s_next, done) tuples (see envs.generate_episode). alpha step size. gamma discount factor. For every step t of the episode, in order t = 0, 1, ..., T-1: V[s_t] <- V[s_t] + alpha * (G_t - V[s_t]), where G_t = r_t + gamma r_{t+1} + gamma^2 r_{t+2} + ... is the return from step t. Compute all the G_t first (a backward pass is easiest: G_t = r_t + gamma G_{t+1}), then apply the updates. If a state appears twice, it is updated twice. Check by hand (Exercise 5.1 in the chapter): V = [1, 2, 0], alpha = 0.5, gamma = 0.9, episode [(0, 1, 1, False), (1, 0, 2, False), (2, 4, None, True)] gives returns G = [4.24, 3.6, 4] and V = [2.62, 2.8, 2.0]. """ # TODO: backward pass for the returns G_t, then the forward pass of updates. raise NotImplementedError("Task 1: mc_update") def td0_update(V: np.ndarray, episode: list[tuple], alpha: float, gamma: float = 1.0) -> None: """Task 2. TD(0) update, applied online, one transition at a time. Inputs as in `mc_update`. For each transition (s, r, s_next, done) in order: delta = r + gamma * V[s_next] - V[s] (use 0 instead of V[s_next] if done) V[s] <- V[s] + alpha * delta "Online" means the update of step t is visible when you process step t+1. Check by hand (Exercise 5.1): same V, alpha, gamma and episode as above give TD errors [1.8, -2, 4] and V = [1.9, 1.0, 2.0]. """ # TODO: one loop over the transitions; remember the terminal case. raise NotImplementedError("Task 2: td0_update") # --------------------------------------------------------------------------- # Part B: control (learning Q and acting on it) # --------------------------------------------------------------------------- def epsilon_greedy(q_values: np.ndarray, epsilon: float, rng: np.random.Generator) -> int: """Task 3. Pick an action epsilon-greedily. Inputs q_values float array of shape (n_actions,): Q(s, .) for the current state. epsilon exploration probability in [0, 1]. rng a numpy Generator; use it for ALL randomness. With probability epsilon return a uniformly random action (any of the n_actions, including the greedy one). Otherwise return an action with the largest Q-value, breaking ties uniformly at random. Returns a Python int. Why random tie-breaking: at the start Q is all zeros, and np.argmax would always pick action 0. Useful: rng.random(), rng.integers(n), rng.choice(array), np.flatnonzero. """ # TODO raise NotImplementedError("Task 3: epsilon_greedy") def sarsa_target(Q: np.ndarray, r: float, s_next: int, a_next: int, done: bool, gamma: float) -> float: """Task 4a. The SARSA target r + gamma * Q(s', a'), where a' is the action the agent WILL actually take next. If done, the target is just r.""" # TODO raise NotImplementedError("Task 4a: sarsa_target") def q_learning_target(Q: np.ndarray, r: float, s_next: int, done: bool, gamma: float) -> float: """Task 4b. The Q-learning target r + gamma * max_a' Q(s', a'). If done, just r.""" # TODO raise NotImplementedError("Task 4b: q_learning_target") def expected_sarsa_target(Q: np.ndarray, r: float, s_next: int, done: bool, gamma: float, epsilon: float) -> float: """Task 4c. The Expected SARSA target r + gamma * sum_a' pi(a'|s') Q(s', a'), where pi is epsilon-greedy with respect to Q. If done, just r. Under epsilon-greedy, every action gets probability epsilon / n_actions, and the greedy action gets an extra (1 - epsilon). (If several actions tie for the max, it does not matter how that extra mass is split: they all have the same value.) """ # TODO raise NotImplementedError("Task 4c: expected_sarsa_target") def sarsa_episode(env, Q: np.ndarray, alpha: float, epsilon: float, gamma: float, rng: np.random.Generator, max_steps: int = 10_000) -> float: """Task 5a. Run ONE episode of SARSA on `env`, updating Q in place. Inputs env has reset() -> s and step(a) -> (s_next, r, done). Q float array (n_states, n_actions), modified in place. alpha, epsilon, gamma step size, exploration rate, discount. rng numpy Generator, passed to epsilon_greedy. max_steps stop the episode after this many steps (a safety net). Returns the episode's total (undiscounted) reward. SARSA, step by step: s = env.reset(); a = epsilon_greedy(Q[s]) repeat: s', r, done = env.step(a) a' = epsilon_greedy(Q[s']) (choose the NEXT action first) Q[s, a] += alpha * (sarsa_target(...) - Q[s, a]) s, a = s', a' until done """ # TODO raise NotImplementedError("Task 5a: sarsa_episode") def q_learning_episode(env, Q: np.ndarray, alpha: float, epsilon: float, gamma: float, rng: np.random.Generator, max_steps: int = 10_000) -> float: """Task 5b. Run ONE episode of Q-learning on `env`, updating Q in place. Same inputs and return value as `sarsa_episode`. Q-learning, step by step: s = env.reset() repeat: a = epsilon_greedy(Q[s]) (behavior policy) s', r, done = env.step(a) Q[s, a] += alpha * (q_learning_target(...) - Q[s, a]) (greedy target) s = s' until done """ # TODO raise NotImplementedError("Task 5b: q_learning_episode")