"""Chapter 8 lab: actor-critic and generalized advantage estimation (GAE). Your file. Fill in the four tasks marked TODO, in order; run the tests after each: python -m pytest -q test_actor_critic.py -k gae # e.g. only the GAE tests Task 1 td_errors one-step TD errors, with episode boundaries Task 2 compute_gae GAE advantages and lambda-returns (backward recursion) Task 3 normalize_advantages zero mean, unit std within a batch Task 4 actor_critic_loss policy loss + value loss - entropy bonus Everything below the line "Provided: A2C on CartPole" is complete scaffolding. """ import numpy as np import torch import torch.nn as nn from cartpole import CartPole torch.set_num_threads(1) # --------------------------------------------------------------------------- # Task 1 # --------------------------------------------------------------------------- def td_errors(rewards, values, last_value, dones, gamma): """One-step TD errors for a rollout of T steps. delta_t = r_t + gamma * V(s_{t+1}) * (1 - done_t) - V(s_t) Args: rewards: array (T,) or (T, N). rewards[t] = r_t, the reward after action a_t. values: array like rewards. values[t] = V(s_t), the critic's value of the state in which a_t was taken. last_value: array () or (N,). V(s_T): the critic's value of the state reached after the last step of the rollout (used to bootstrap). dones: array like rewards, 0/1. dones[t] = 1 if the episode ENDED at step t, i.e. s_{t+1} is terminal (worth 0). The next row of the rollout then belongs to a new episode, so values[t + 1] must NOT be used. gamma: discount factor. The tests pass plain Python lists and floats too (e.g. last_value=0.0), so convert every input with np.asarray(..., dtype=np.float64) first. Returns: deltas: float array with the same shape as rewards. """ # TODO (Task 1): about 5 lines. # - Convert rewards, values, last_value, dones with np.asarray(..., dtype=np.float64). # - Build next_values: values shifted up by one row, with last_value as the final row # (np.concatenate along axis 0; last_value[None] adds the missing leading axis). # - Zero the bootstrap where dones[t] = 1. raise NotImplementedError("Task 1: td_errors") # --------------------------------------------------------------------------- # Task 2 # --------------------------------------------------------------------------- def compute_gae(rewards, values, last_value, dones, gamma, lam): """Generalized advantage estimation, computed backward in time. A_t = delta_t + gamma * lam * (1 - done_t) * A_{t+1}, with A_T = 0, which equals sum_l (gamma*lam)^l delta_{t+l}, stopping at the end of the episode. lam = 0 gives the one-step TD error; lam = 1 gives (bootstrapped) Monte Carlo return minus V(s_t). Args: as in td_errors, plus lam: the GAE parameter lambda in [0, 1]. Returns: advantages: float array, same shape as rewards. returns: advantages + values. These are the lambda-returns, the regression targets for the critic. """ # TODO (Task 2): about 8 lines. # - deltas = td_errors(...); convert dones and values with np.asarray as in Task 1. # - Loop t = T-1, ..., 0 carrying "next_adv" (start at zeros_like(deltas[0])). # Reset the carry at episode ends with (1 - dones[t]). # - returns = advantages + values raise NotImplementedError("Task 2: compute_gae") # --------------------------------------------------------------------------- # Task 3 # --------------------------------------------------------------------------- def normalize_advantages(adv, eps=1e-8): """Shift and scale a batch of advantages to mean 0 and standard deviation 1. A_hat <- (A_hat - mean(A_hat)) / (std(A_hat) + eps) Use the population standard deviation (divide by n, not n - 1). Args: adv: torch tensor of any shape (all entries form one batch). Returns: torch tensor, same shape. """ # TODO (Task 3): one line. Hint: tensor.std(unbiased=False). raise NotImplementedError("Task 3: normalize_advantages") # --------------------------------------------------------------------------- # Task 4 # --------------------------------------------------------------------------- def actor_critic_loss(logits, values, actions, advantages, returns, value_coef=0.5, entropy_coef=0.01): """The A2C loss for one batch of B transitions. policy_loss = -mean_b[ log pi(a_b | s_b) * A_b ] value_loss = mean_b[ (V(s_b) - R_b)^2 ] entropy = mean_b[ H(pi(. | s_b)) ], H(p) = -sum_a p_a log p_a total = policy_loss + value_coef * value_loss - entropy_coef * entropy Advantages and returns are targets: treat them as constants (no gradient flows into them), even if the caller forgot to detach them. Args: logits: (B, n_actions) float tensor, the actor's output. values: (B,) float tensor, the critic's V(s_b). actions: (B,) long tensor, the actions taken. advantages: (B,) float tensor, A_b (already normalized, if you normalize). returns: (B,) float tensor, R_b, the critic's regression targets. Returns: total: scalar tensor (differentiable w.r.t. logits and values). info: dict with python floats "policy_loss", "value_loss", "entropy". """ # TODO (Task 4): about 9 lines. # - Detach advantages and returns. # - log_probs = torch.log_softmax(logits, dim=-1); pick the taken actions with .gather. # - Entropy per state: -(p * log p).sum(-1), with p = log_probs.exp(). # - Return (total, {"policy_loss": ..., "value_loss": ..., "entropy": ...}) with .item() floats. raise NotImplementedError("Task 4: actor_critic_loss") # =========================================================================== # Provided: A2C on CartPole. Nothing to write below; read it to see how the # four functions above fit together. # =========================================================================== class ActorCritic(nn.Module): """Actor (policy logits) and critic (state value) for a discrete-action task. shared=False: two separate MLPs. shared=True: one trunk with two heads. """ def __init__(self, obs_dim=4, n_actions=2, hidden=64, shared=False): super().__init__() self.shared = shared if shared: self.trunk = nn.Sequential(nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh()) self.pi_head = nn.Linear(hidden, n_actions) self.v_head = nn.Linear(hidden, 1) else: self.actor = nn.Sequential( nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, n_actions) ) self.critic = nn.Sequential( nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1) ) def forward(self, obs): """obs (B, obs_dim) -> logits (B, n_actions), value (B,).""" if self.shared: h = self.trunk(obs) return self.pi_head(h), self.v_head(h).squeeze(-1) return self.actor(obs), self.critic(obs).squeeze(-1) class VecCartPole: """N independent CartPoles stepped together, with automatic reset. step() reports `terminated` (the pole fell or the cart left the track: a true terminal state, worth 0) separately from `truncated` (the 500-step time limit: the state is NOT terminal, we just stopped watching). """ def __init__(self, n_envs, seed=0, max_steps=500): self.envs = [CartPole(seed=seed + 1000 * i, max_steps=max_steps) for i in range(n_envs)] self.ep_returns = np.zeros(n_envs) self.finished = [] # returns of completed episodes, in order def reset(self): return np.stack([e.reset() for e in self.envs]) def step(self, actions): n = len(self.envs) obs = np.zeros((n, 4), dtype=np.float32) final_obs = np.zeros((n, 4), dtype=np.float32) rewards = np.zeros(n) terminated = np.zeros(n, dtype=bool) truncated = np.zeros(n, dtype=bool) for i, (env, a) in enumerate(zip(self.envs, actions)): o, r, done = env.step(int(a)) rewards[i] = r final_obs[i] = o self.ep_returns[i] += r if done: truncated[i] = env.t >= env.max_steps and abs(env.state[0]) <= env.x_threshold and abs( env.state[2] ) <= env.theta_threshold terminated[i] = not truncated[i] self.finished.append(self.ep_returns[i]) self.ep_returns[i] = 0.0 o = env.reset() obs[i] = o return obs, rewards, terminated, truncated, final_obs def collect_rollout(model, envs, obs, n_steps, gamma): """Run the current policy for n_steps in every env. Returns (batch, next_obs). batch["obs"] (n_steps, N, 4), ["actions"] (n_steps, N), ["rewards"], ["values"], ["dones"] (n_steps, N), and ["last_value"] (N,). Time-limit truncation: the episode is cut off but s_{t+1} is not terminal. We fold the missing future into the reward, r_t <- r_t + gamma * V(s_{t+1}), and then mark the step done so GAE does not leak into the next episode. (Exercise 8.9.) """ n = obs.shape[0] B = {k: np.zeros((n_steps, n), dtype=np.float32) for k in ("actions", "rewards", "values", "dones")} B["obs"] = np.zeros((n_steps, n, 4), dtype=np.float32) for t in range(n_steps): with torch.no_grad(): logits, v = model(torch.as_tensor(obs)) a = torch.distributions.Categorical(logits=logits).sample().numpy() next_obs, r, term, trunc, final_obs = envs.step(a) if trunc.any(): with torch.no_grad(): _, v_final = model(torch.as_tensor(final_obs)) r = r + gamma * v_final.numpy() * trunc B["obs"][t], B["actions"][t], B["rewards"][t] = obs, a, r B["values"][t], B["dones"][t] = v.numpy(), (term | trunc) obs = next_obs with torch.no_grad(): B["last_value"] = model(torch.as_tensor(obs))[1].numpy() return B, obs def train_a2c( total_steps=100_000, n_envs=8, n_steps=32, gamma=0.99, lam=0.95, lr=1e-3, value_coef=0.5, entropy_coef=0.01, normalize=True, shared=False, max_grad_norm=0.5, seed=0, log_every=5_000, verbose=False, ): """Synchronous advantage actor-critic (A2C) on CartPole. Each update: collect n_steps x n_envs transitions with the current policy, compute GAE advantages and lambda-returns, take ONE gradient step on actor_critic_loss. Returns a dict with "steps" and "avg_return" (mean return of the last 20 finished episodes, logged every log_every env steps), "episode_returns" (all of them) and the trained "model". """ torch.manual_seed(seed) np.random.seed(seed) model = ActorCritic(shared=shared) opt = torch.optim.Adam(model.parameters(), lr=lr) envs = VecCartPole(n_envs, seed=seed) obs = envs.reset() steps, next_log = 0, log_every hist = {"steps": [], "avg_return": []} while steps < total_steps: batch, obs = collect_rollout(model, envs, obs, n_steps, gamma) steps += n_steps * n_envs adv, ret = compute_gae(batch["rewards"], batch["values"], batch["last_value"], batch["dones"], gamma, lam) logits, values = model(torch.as_tensor(batch["obs"]).reshape(-1, 4)) adv_t = torch.as_tensor(adv, dtype=torch.float32).reshape(-1) if normalize: adv_t = normalize_advantages(adv_t) loss, info = actor_critic_loss( logits, values, torch.as_tensor(batch["actions"]).reshape(-1).long(), adv_t, torch.as_tensor(ret, dtype=torch.float32).reshape(-1), value_coef=value_coef, entropy_coef=entropy_coef, ) opt.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) opt.step() if steps >= next_log: next_log += log_every recent = envs.finished[-20:] avg = float(np.mean(recent)) if recent else 0.0 hist["steps"].append(steps) hist["avg_return"].append(avg) if verbose: print( f"steps {steps:7d} avg return {avg:6.1f} entropy {info['entropy']:.3f} " f"value loss {info['value_loss']:.2f}" ) hist["episode_returns"] = list(envs.finished) hist["model"] = model return hist