Exercise 04: Actor-critic on CartPole

"""Actor-critic on CartPole-v1 — exercise starter.

Fill in the TODOs, then from the repo root:

    pip install -r exercises/requirements.txt
    pytest exercises/04-actor-critic/

See README.md for the task and HINTS.md if you're stuck. The reference
implementation is in solution/actor_critic.py — look after you've tried.
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

GAMMA = 0.99


class PolicyNet(nn.Module):
    """Maps a state vector to action logits (raw scores, NOT probabilities)."""

    def __init__(self, state_dim: int, n_actions: int, hidden: int = 128):
        super().__init__()
        # TODO: build an MLP  state_dim -> hidden -> hidden -> n_actions  with
        # ReLU between the linear layers. Same shape as exercise 02 — no softmax,
        # return logits. (Categorical(logits=...) handles the softmax internally
        # and is numerically more stable.)
        raise NotImplementedError("PolicyNet.__init__")

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        # TODO: return logits, shape [..., n_actions]
        raise NotImplementedError("PolicyNet.forward")


class ValueNet(nn.Module):
    """Maps a state vector to a scalar estimate of V(s)."""

    def __init__(self, state_dim: int, hidden: int = 128):
        super().__init__()
        # TODO: build an MLP  state_dim -> hidden -> hidden -> 1  with ReLU
        # between the linear layers. The output is a single scalar V(s) —
        # squeeze the trailing dimension in forward so the output shape is
        # [...] (not [..., 1]).
        raise NotImplementedError("ValueNet.__init__")

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        # TODO: return V(s) — shape [] for a single state, [batch] for a batch.
        # Hint: nn.Linear(..., 1) gives shape [..., 1]; use .squeeze(-1) to drop
        # the last dim.
        raise NotImplementedError("ValueNet.forward")


def select_action(policy: PolicyNet, state: np.ndarray):
    """Sample an action from the policy.

    Returns (action, log_prob):
      - action: a Python int
      - log_prob: a scalar tensor for log π(action | state) that STILL CARRIES
        GRADIENT (don't call .item() on it — the loss has to backprop through it).
    """
    # TODO: same as exercise 02 —
    #   - turn `state` into a float32 tensor
    #   - logits = policy(state_tensor)
    #   - dist = torch.distributions.Categorical(logits=logits)
    #   - action = dist.sample()
    #   - return int(action.item()), dist.log_prob(action)
    raise NotImplementedError("select_action")


def compute_returns(rewards: list, gamma: float = GAMMA) -> torch.Tensor:
    """Discounted return-to-go for each timestep.

        returns[t] = rewards[t] + gamma*rewards[t+1] + gamma**2*rewards[t+2] + ...

    Return a 1-D float32 tensor the same length as `rewards`. Same function as
    exercise 02 — copy it if you like, but re-derive it; you'll need the
    intuition when you get to TD methods.
    """
    # TODO: walk `rewards` backwards, accumulating  R = r + gamma * R,
    # building the list of returns, then reverse it and make it a tensor.
    raise NotImplementedError("compute_returns")


def actor_critic_loss(log_probs: list, returns: torch.Tensor, values: torch.Tensor):
    """Compute actor and critic losses separately.

    The key difference from REINFORCE: instead of normalizing the raw returns
    as a crude baseline, we use the critic's estimate V(s_t) as a state-dependent
    baseline. The advantage A_t = G_t - V(s_t) tells the actor whether the
    outcome was better or worse than the critic predicted.

    Important: the policy gradient does NOT flow through the critic. The
    advantages must be detached from the critic's computation graph, or the
    actor loss will pull ValueNet's parameters in a policy direction, which
    is wrong. Use `(returns - values).detach()`.

    Returns:
        actor_loss: scalar tensor. Requires grad (flows into PolicyNet).
        critic_loss: scalar tensor. Requires grad (flows into ValueNet).
    """
    # TODO:
    #   - advantages = (returns - values).detach()     # stop gradient here
    #   - lp = torch.stack(list(log_probs))
    #   - actor_loss  = -(lp * advantages).sum()       # policy gradient, negated to minimize
    #   - critic_loss = F.mse_loss(values, returns)    # train the critic with MSE
    #   - return actor_loss, critic_loss
    raise NotImplementedError("actor_critic_loss")


def train(
    num_episodes: int = 600,
    lr: float = 3e-3,
    seed: int = 0,
    gamma: float = GAMMA,
    value_coef: float = 0.5,
) -> list:
    """Train actor-critic on CartPole-v1. Returns the list of per-episode total rewards."""
    import gymnasium as gym

    torch.manual_seed(seed)
    np.random.seed(seed)
    env = gym.make("CartPole-v1")
    state_dim = env.observation_space.shape[0]
    n_actions = env.action_space.n

    policy = PolicyNet(state_dim, n_actions)
    critic = ValueNet(state_dim)

    # One Adam optimizer over both networks — simpler than two separate ones.
    optimizer = torch.optim.Adam(
        list(policy.parameters()) + list(critic.parameters()), lr=lr
    )

    episode_returns = []
    for ep in range(num_episodes):
        state, _ = env.reset(seed=seed + ep)
        log_probs, rewards, values = [], [], []
        done = False
        while not done:
            state_t = torch.as_tensor(state, dtype=torch.float32)
            action, log_prob = select_action(policy, state)
            # TODO: compute value = critic(state_t) and append it to `values`.
            # This is V(s_t) before we take the action — the baseline for this step.
            raise NotImplementedError("train: call critic(state_t) and append value")

            state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            log_probs.append(log_prob)
            rewards.append(float(reward))

        # TODO: one actor-critic update for this episode:
        #   returns = compute_returns(rewards, gamma)
        #   values_t = torch.stack(values)
        #   actor_loss, critic_loss = actor_critic_loss(log_probs, returns, values_t)
        #   loss = actor_loss + value_coef * critic_loss
        #   optimizer.zero_grad()
        #   loss.backward()
        #   torch.nn.utils.clip_grad_norm_(
        #       list(policy.parameters()) + list(critic.parameters()), 1.0
        #   )
        #   optimizer.step()
        raise NotImplementedError("train: the update step")

        episode_returns.append(sum(rewards))

    env.close()
    return episode_returns


if __name__ == "__main__":
    rets = train()
    last100 = rets[-100:]
    print(
        f"episodes: {len(rets)}  "
        f"mean(last 100): {sum(last100) / len(last100):.1f}  "
        f"max: {max(rets):.0f}  "
        f"(random policy ≈ 21, solved ≈ 475)"
    )

Both are plain MLPs. For PolicyNet, it’s the same nn.Sequential(Linear, ReLU, Linear, ReLU, Linear) as exercise 02: state_dim → hidden → hidden → n_actions, no softmax. For ValueNet, same structure but the last layer outputs 1: state_dim → hidden → hidden → 1. In ValueNet.forward, call .squeeze(-1) on the output so the shape is [] for a single state and [batch] for a batch, not [1] or [batch, 1]. The test checks the shape exactly.

These are identical to exercise 02. Copy them if you like, but re-read them, compute_returns in particular. Walking backwards (start R = 0, then R = r + gamma * R for each r in reversed(rewards)) gives you G_t because each step incorporates all future rewards. When you do it forward, you’d need two nested loops. Reverse is O(n) and clean.

This is the only structural difference from REINFORCE. You need:

advantages = (returns - values).detach()

The .detach() cuts the computation graph so that when actor_loss.backward() runs, it produces gradients in PolicyNet (via log_probs) but not in ValueNet (via values). Without it, the actor loss would try to pull ValueNet’s parameters in the direction that makes advantages look bigger. That’s not training the critic to predict returns, that’s just breaking it. There’s a test that checks this: test_advantages_do_not_flow_into_value_net.

After the detach:

lp = torch.stack(list(log_probs))
actor_loss  = -(lp * advantages).sum()
critic_loss = F.mse_loss(values, returns)
return actor_loss, critic_loss

The minus sign on actor_loss is the same reason as in REINFORCE: the optimizer minimizes, but the policy gradient is an ascent direction. F.mse_loss(values, returns) trains the critic to predict G_t. Returns here are the Monte Carlo targets and the critic is learning to match them. Note that critic_loss uses values without detaching; it needs gradient to flow into ValueNet.

The loop structure is given; you need to add one line inside while not done:

state_t = torch.as_tensor(state, dtype=torch.float32)   # already there
action, log_prob = select_action(policy, state)           # already there
value = critic(state_t)                                   # <-- add this
# ... env.step, append log_prob and reward ...
values.append(value)                                      # <-- and this

Call critic before env.step; you want V(s_t) at the state before the action, not the next state. Keep the tensors in the list (don’t .item() them); you’ll torch.stack(values) after the episode to get one tensor the critic loss can differentiate through.

After the episode loop:

returns = compute_returns(rewards, gamma)
values_t = torch.stack(values)
actor_loss, critic_loss = actor_critic_loss(log_probs, returns, values_t)
loss = actor_loss + value_coef * critic_loss

optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(
    list(policy.parameters()) + list(critic.parameters()), 1.0
)
optimizer.step()

value_coef=0.5 scales the critic’s contribution: without it the critic loss (which can be in the hundreds for early episodes, since V(s) starts near 0 and returns are ~10–500) would dominate and destabilize the actor update. The gradient clip (norm at 1.0) caps any single episode’s impact on the weights.

If you’re past hint 6 and still stuck, read solution/actor_critic.py. Then come back and re-derive it without looking; you don’t own it until you can.