"""Chapter 7 lab: policy gradients and REINFORCE (your file: fill in the TODOs). Part A (numpy): a softmax policy on a K-armed bandit. You compute the exact policy gradient, the single-sample REINFORCE estimates, and check that they agree on average and that a baseline shrinks their spread. Part B (PyTorch): REINFORCE on CartPole, with three choices of per-step weight: "total" every action is weighted by the whole-episode return G_0 "to_go" each action is weighted by its reward-to-go G_t "baseline" each action is weighted by G_t - V(s_t), V a learned baseline Notation follows the chapter: theta are logits, pi = softmax(theta), J(theta) = sum_a pi(a) q(a), G_t = sum_k gamma^k r_{t+k}. """ import numpy as np import torch import torch.nn as nn from cartpole import CartPole torch.set_num_threads(1) # --------------------------------------------------------------------------- # Part A: the softmax bandit (numpy) # --------------------------------------------------------------------------- def softmax(theta): """Softmax of a 1-D array of logits, computed stably. (Provided.)""" z = np.asarray(theta, dtype=float) z = z - z.max() e = np.exp(z) return e / e.sum() def softmax_score(theta, actions): """Task 1. The score: gradient of log pi_theta(a) w.r.t. the logits. Args: theta: (K,) logits. actions: (N,) integer actions in 0..K-1. Returns: (N, K) array; row i is grad_theta log pi_theta(actions[i]). Formula (chapter, "the score of a softmax policy"): d log pi(a) / d theta_k = 1[a == k] - pi_k, i.e. e_a - pi. """ # TODO: build one-hot rows for `actions` (np.eye(K)[actions] works) and subtract pi from each row. raise NotImplementedError("Task: softmax_score") def exact_gradient(theta, q): """Task 2. The exact policy gradient of a softmax bandit. Args: theta: (K,) logits. q: (K,) true mean reward of each arm. Returns: (K,) array grad_theta J, where J(theta) = sum_a pi_a q_a. Formula (chapter, worked example): dJ / d theta_k = pi_k * (q_k - J). """ # TODO: compute pi and J(theta), then apply the formula above (one line, vectorised). raise NotImplementedError("Task: exact_gradient") def reinforce_estimates(theta, actions, rewards, baseline=0.0): """Task 3. Single-sample REINFORCE gradient estimates. Args: theta: (K,) logits. actions: (N,) actions sampled from pi_theta. rewards: (N,) rewards received for those actions. baseline: a number b (must not depend on the sampled action). Returns: (N, K) array; row i is (rewards[i] - b) * grad_theta log pi_theta(actions[i]). Its mean over rows is the usual N-sample REINFORCE estimate. """ # TODO: use softmax_score, then scale row i by (rewards[i] - baseline). raise NotImplementedError("Task: reinforce_estimates") def sample_bandit(theta, q, n, rng, noise=1.0): """Draw n (action, reward) pairs from a Gaussian bandit. (Provided.) Actions come from softmax(theta); reward = q[a] + noise * standard normal. """ pi = softmax(theta) actions = rng.choice(len(pi), size=n, p=pi) rewards = np.asarray(q, dtype=float)[actions] + noise * rng.standard_normal(n) return actions, rewards def total_variance(estimates): """Trace of the covariance of the rows of `estimates`: E||g - mean g||^2. (Provided.)""" est = np.asarray(estimates, dtype=float) return float(((est - est.mean(axis=0)) ** 2).sum(axis=1).mean()) # --------------------------------------------------------------------------- # Part B: REINFORCE on CartPole (PyTorch) # --------------------------------------------------------------------------- def rewards_to_go(rewards, gamma): """Task 4. Discounted reward-to-go for one episode. Args: rewards: (T,) rewards r_0 .. r_{T-1}. gamma: discount factor. Returns: (T,) float array with G_t = r_t + gamma * r_{t+1} + gamma^2 * r_{t+2} + ... Compute it with one backward pass: G_t = r_t + gamma * G_{t+1}, G_T = 0. (Like almost every implementation, we drop the extra gamma^t factor that the exact gradient of the discounted objective would put in front of G_t; see the chapter's note on this.) """ # TODO: loop t = T-1 down to 0, keeping a running sum running = r_t + gamma * running. raise NotImplementedError("Task: rewards_to_go") def reinforce_loss(logits, actions, weights): """Task 5. The surrogate loss whose gradient is the REINFORCE estimate. Args: logits: (N, A) tensor of action logits for N visited states (requires grad). actions: (N,) long tensor of the actions that were taken. weights: (N,) tensor, the weight of each action (G_0, G_t or G_t - b). Treat these as constants: no gradient may flow into them. Returns: Scalar tensor L = -(1/N) * sum_i weights[i] * log pi(actions[i] | s_i). Minimising L with gradient descent = gradient ascent on J. """ # TODO: log-probs via torch.log_softmax(logits, dim=-1); pick the taken actions with .gather; # multiply by weights.detach(); average; negate. raise NotImplementedError("Task: reinforce_loss") class PolicyNet(nn.Module): """Small MLP: state (4,) -> logits for the 2 actions. (Provided.)""" def __init__(self, obs_dim=4, n_actions=2, hidden=32): super().__init__() self.net = nn.Sequential( nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, n_actions), ) def forward(self, s): return self.net(s) class ValueNet(nn.Module): """Small MLP: state (4,) -> scalar baseline V(s). (Provided.)""" def __init__(self, obs_dim=4, hidden=32): super().__init__() self.net = nn.Sequential( nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, 1), ) def forward(self, s): return self.net(s).squeeze(-1) def run_episode(env, policy, rng): """Play one episode with the stochastic policy. (Provided.) `rng` is a numpy Generator used to sample actions from pi(. | s). Returns (states (T, 4) float32, actions (T,) int64, rewards (T,) float64). """ states, actions, rewards = [], [], [] s = env.reset() done = False while not done: with torch.no_grad(): probs = torch.softmax(policy(torch.as_tensor(s)), dim=-1).numpy() a = int(np.searchsorted(np.cumsum(probs), rng.random() * probs.sum())) a = min(a, len(probs) - 1) s_next, r, done = env.step(a) states.append(s) actions.append(a) rewards.append(r) s = s_next return np.array(states, dtype=np.float32), np.array(actions), np.array(rewards, dtype=float) def train_reinforce(variant="baseline", n_iters=50, episodes_per_iter=8, gamma=0.99, lr=0.02, value_lr=0.02, value_steps=20, seed=0, log_every=0): """REINFORCE on CartPole. (Provided scaffolding; it calls your Tasks 4 and 5.) Each iteration: play `episodes_per_iter` episodes with the current policy, turn their rewards into per-step weights according to `variant`, take one gradient step on reinforce_loss, then (for "baseline") refit the value net. Returns a 1-D array with the (undiscounted) return of every episode played. """ assert variant in ("total", "to_go", "baseline"), variant torch.manual_seed(seed) env = CartPole(seed=seed) rng = np.random.default_rng(seed) policy = PolicyNet() value = ValueNet() opt = torch.optim.Adam(policy.parameters(), lr=lr) opt_v = torch.optim.Adam(value.parameters(), lr=value_lr) history = [] for it in range(n_iters): S, A, W = [], [], [] for _ in range(episodes_per_iter): s, a, r = run_episode(env, policy, rng) history.append(r.sum()) g = rewards_to_go(r, gamma) if variant == "total": w = np.full_like(g, g[0]) # every action gets the whole return else: w = g # reward-to-go S.append(s) A.append(a) W.append(w) S = torch.as_tensor(np.concatenate(S)) A = torch.as_tensor(np.concatenate(A)) G = torch.as_tensor(np.concatenate(W), dtype=torch.float32) if variant == "baseline": with torch.no_grad(): weights = G - value(S) # advantage estimate G_t - V(s_t) else: weights = G loss = reinforce_loss(policy(S), A, weights) opt.zero_grad() loss.backward() opt.step() if variant == "baseline": # regress V(s_t) toward the Monte Carlo return G_t for _ in range(value_steps): v_loss = ((value(S) - G) ** 2).mean() opt_v.zero_grad() v_loss.backward() opt_v.step() if log_every and (it + 1) % log_every == 0: recent = np.mean(history[-episodes_per_iter:]) print(f"[{variant}] iter {it + 1:3d} mean return {recent:6.1f}") return np.array(history)