"""Chapter 6 lab: function approximation and DQN. Your file. Fill in the five tasks in order; each one has tests in test_dqn.py. Task 1 linear_td0_update semi-gradient TD(0) with linear features (numpy) Task 2 ReplayBuffer push / sample a circular buffer of transitions Task 3 QNetwork, epsilon_by_step, select_action the acting side Task 4 compute_td_targets, dqn_loss the learning signal Task 5 update_target, train_step one DQN update Run the tests for one task at a time, e.g. python -m pytest -q -k task1 The training loop that uses all of this is in train.py (provided, worth reading). `python run.py` trains on CartPole once the tests pass. """ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F torch.set_num_threads(1) # ============================================================ Task 1 def linear_td0_update(w, phi, r, phi_next, done, alpha, gamma): """One semi-gradient TD(0) update of a *linear* value function. The value estimate is V(s; w) = w @ phi(s). The update is delta = r + gamma * V(s'; w) - V(s; w) (V(s'; w) := 0 if done) w <- w + alpha * delta * phi(s) "Semi-gradient": the target r + gamma * V(s'; w) also depends on w, but we treat it as a constant, so the gradient is just grad_w V(s; w) = phi(s). Args: w: (d,) float array, current weights. Do not modify it in place. phi: (d,) features of the current state s. r: float, reward received after leaving s. phi_next: (d,) features of the next state s'. done: bool, True if s' is terminal (then V(s') = 0). alpha: step size. gamma: discount factor. Returns: (d,) float array: the new weights (a new array). """ # TODO: compute V(s), V(s') (0 if done), the TD error delta, and return # w + alpha * delta * phi. About 4 lines. raise NotImplementedError("Task 1: linear_td0_update") # ============================================================ Task 2 class ReplayBuffer: """A fixed-size circular buffer of transitions (s, a, r, s', done). When full, a new transition overwrites the *oldest* one. `sample` draws a minibatch uniformly at random (with replacement) from what is stored. The storage arrays are allocated for you in __init__. You write `push` and `sample`. """ def __init__(self, capacity, obs_dim, seed=0): self.capacity = int(capacity) self.obs = np.zeros((capacity, obs_dim), dtype=np.float32) self.actions = np.zeros(capacity, dtype=np.int64) self.rewards = np.zeros(capacity, dtype=np.float32) self.next_obs = np.zeros((capacity, obs_dim), dtype=np.float32) self.dones = np.zeros(capacity, dtype=np.float32) self.rng = np.random.default_rng(seed) self.idx = 0 # slot where the next transition will be written self.size = 0 # number of transitions currently stored (<= capacity) def push(self, obs, action, reward, next_obs, done): """Store one transition in slot `self.idx`, then advance `idx` (wrapping around to 0 at `capacity`) and grow `size` (up to capacity). `done` should be stored as 1.0 (True) or 0.0 (False). """ # TODO: write the five fields into slot self.idx, then update # self.idx (use % self.capacity) and self.size (use min). raise NotImplementedError("Task 2: ReplayBuffer.push") def sample(self, batch_size): """Sample `batch_size` stored transitions uniformly at random, with replacement. Use `self.rng`. Returns a dict of torch tensors: "obs": (B, obs_dim) float32 "actions": (B,) int64 "rewards": (B,) float32 "next_obs": (B, obs_dim) float32 "dones": (B,) float32 (1.0 where s' is terminal) """ # TODO: draw indices with self.rng.integers(0, self.size, size=batch_size) # (never index past self.size: those slots are still empty), then # build the dict with torch.from_numpy(self.obs[ix]) etc. raise NotImplementedError("Task 2: ReplayBuffer.sample") def __len__(self): return self.size # ============================================================ Task 3 class QNetwork(nn.Module): """An MLP mapping a state to one Q-value per action. Architecture: Linear(obs_dim, hidden) -> ReLU -> Linear(hidden, hidden) -> ReLU -> Linear(hidden, n_actions) forward(obs): (B, obs_dim) float tensor -> (B, n_actions) Q-values. One forward pass gives the values of *all* actions, so max_a Q(s, a) costs a single pass (DQN's design choice). """ def __init__(self, obs_dim, n_actions, hidden=64): super().__init__() # TODO: build self.net = nn.Sequential(...) with the layers above. raise NotImplementedError("Task 3: QNetwork.__init__") def forward(self, obs): # TODO: return self.net(obs) raise NotImplementedError("Task 3: QNetwork.forward") def epsilon_by_step(step, eps_start=1.0, eps_end=0.05, decay_steps=10_000): """Linearly anneal epsilon from eps_start (at step 0) to eps_end (at step decay_steps), then hold it at eps_end. eps(step) = eps_start + (eps_end - eps_start) * min(step / decay_steps, 1) """ # TODO: one or two lines. raise NotImplementedError("Task 3: epsilon_by_step") def select_action(q_net, obs, epsilon, rng): """Epsilon-greedy action from a Q-network. With probability epsilon pick a uniformly random action, otherwise the greedy action argmax_a Q(obs, a). Use `rng.random()` for the coin flip and `rng.integers(n_actions)` for the random action. Args: q_net: a QNetwork (or any module mapping (1, obs_dim) -> (1, n_actions)). obs: (obs_dim,) numpy array, a single state. epsilon: float in [0, 1]. rng: numpy Generator. Returns: a Python int. """ # TODO: # 1. under torch.no_grad(), q = q_net(obs as a (1, obs_dim) float32 tensor) # 2. if rng.random() < epsilon: return a random int in [0, n_actions) # 3. else return int(argmax of q) raise NotImplementedError("Task 3: select_action") # ============================================================ Task 4 def compute_td_targets(target_net, rewards, next_obs, dones, gamma, online_net=None): """The DQN regression targets y, computed WITHOUT gradient. Standard DQN (online_net is None): y = r + gamma * (1 - done) * max_a' Q(s', a'; theta_minus) Double DQN (online_net given): choose a' with the online net, evaluate it with the target net: a* = argmax_a' Q(s', a'; theta) y = r + gamma * (1 - done) * Q(s', a*; theta_minus) Args: target_net: network with parameters theta_minus. rewards: (B,) float tensor. next_obs: (B, obs_dim) float tensor. dones: (B,) float tensor, 1.0 where s' is terminal. gamma: discount factor. online_net: optional network with parameters theta (for Double DQN). Returns: (B,) float tensor that does not require grad. """ # TODO: everything inside `with torch.no_grad():`. # Useful: tensor.max(dim=1).values, tensor.argmax(dim=1, keepdim=True), # tensor.gather(1, index).squeeze(1) raise NotImplementedError("Task 4: compute_td_targets") def dqn_loss(q_net, target_net, batch, gamma, double=False): """The DQN loss on a minibatch: mean Huber loss between Q(s, a; theta) and the (constant) targets y from `compute_td_targets`. loss = mean_i Huber( Q(s_i, a_i; theta) - y_i ), Huber(d) = 0.5 d^2 if |d| <= 1 else |d| - 0.5 Args: q_net: online network (theta). Gradients flow only into this one. target_net: target network (theta_minus). batch: dict from ReplayBuffer.sample. gamma: discount factor. double: if True, use Double DQN targets. Returns: scalar tensor (requires grad through q_net only). """ # TODO: # 1. q_sa = Q(s_i, a_i; theta): run q_net on batch["obs"] -> (B, n_actions), # then pick column a_i in each row with .gather(1, actions.unsqueeze(1)) # 2. y = compute_td_targets(..., online_net=q_net if double else None) # 3. return F.smooth_l1_loss(q_sa, y) (this IS the Huber loss above) raise NotImplementedError("Task 4: dqn_loss") # ============================================================ Task 5 def update_target(target_net, q_net, tau=1.0): """Move the target network's parameters toward the online network's: theta_minus <- tau * theta + (1 - tau) * theta_minus tau = 1.0 is DQN's hard copy (done every C steps); a small tau like 0.005 applied every step is a "soft" (Polyak) update. No gradients involved. """ # TODO: under torch.no_grad(), loop over # zip(target_net.parameters(), q_net.parameters()) # and update each target parameter in place (p_t.mul_(...).add_(...), # or p_t.copy_(...)). raise NotImplementedError("Task 5: update_target") def train_step(q_net, target_net, optimizer, buffer, batch_size, gamma, double=False, max_grad_norm=10.0): """One gradient step of DQN: 1. sample a minibatch from the buffer 2. compute the loss with `dqn_loss` 3. optimizer.zero_grad(); loss.backward() 4. clip the gradient norm to max_grad_norm (torch.nn.utils.clip_grad_norm_) 5. optimizer.step() Does NOT touch target_net (the training loop syncs it separately). Returns: the loss as a Python float. """ # TODO: the five steps above, then return float(loss.item()). raise NotImplementedError("Task 5: train_step")