Actor-critic methods

Setup#

pip install gymnasium torch numpy

CPU only. This runs in about twenty-five seconds — roughly half the time REINFORCE took to reach the same score.

Actor and critic, both from scratch#

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(0); np.random.seed(0)
torch.set_num_threads(1)

env = gym.make("CartPole-v1")

actor  = nn.Sequential(nn.Linear(4, 64), nn.ReLU(), nn.Linear(64, 2))   # what to do
critic = nn.Sequential(nn.Linear(4, 64), nn.ReLU(), nn.Linear(64, 1))   # how good this state is
opt = torch.optim.Adam(list(actor.parameters()) + list(critic.parameters()), lr=0.01)
GAMMA = 0.99
scores = []

for episode in range(600):
    state, _ = env.reset(seed=episode)
    states, actions, rewards, done = [], [], [], False

    while not done:
        logits = actor(torch.tensor(state))
        dist = torch.distributions.Categorical(logits=logits)
        action = dist.sample()
        states.append(state); actions.append(int(action))
        state, reward, terminated, truncated, _ = env.step(int(action))
        rewards.append(reward)
        done = terminated or truncated

    returns, running = [], 0.0
    for r in reversed(rewards):
        running = r + GAMMA * running
        returns.insert(0, running)

    s = torch.tensor(np.array(states))
    a = torch.tensor(actions)
    g = torch.tensor(returns, dtype=torch.float32)

    values = critic(s).squeeze(1)
    advantage = g - values                     # was the outcome better than the critic expected?

    dist = torch.distributions.Categorical(logits=actor(s))
    actor_loss  = -(dist.log_prob(a) * advantage.detach()).mean()
    critic_loss = nn.functional.mse_loss(values, g)
    entropy     = dist.entropy().mean()        # a nudge to keep some randomness alive

    loss = actor_loss + 0.5 * critic_loss - 0.01 * entropy
    opt.zero_grad(); loss.backward(); opt.step()

    scores.append(sum(rewards))
    if (episode + 1) % 100 == 0:
        print(f"episode {episode + 1:3d}  avg of last 100: {np.mean(scores[-100:]):6.1f}"
              f"   critic's guess at the start state: {critic(torch.tensor(states[0])).item():6.1f}")

env.close()

episode 100  avg of last 100:   15.3   critic's guess at the start state:   15.0
episode 200  avg of last 100:  271.5   critic's guess at the start state:   73.3
episode 300  avg of last 100:  353.5   critic's guess at the start state:   67.2
episode 400  avg of last 100:  236.8   critic's guess at the start state:   84.9
episode 500  avg of last 100:  443.5   critic's guess at the start state:   92.7
episode 600  avg of last 100:  500.0   critic's guess at the start state:  104.6

Exact numbers vary by PyTorch build and machine even with seeds fixed. The trend is what matters.

Watch the critic learn its own job#

The second column is the interesting one. The critic's estimate of the starting position climbs: 15, then 73, then 92, then 105.

Early on it says 15 because the pole really does fall after about 15 steps. It is not wrong; the agent is bad. As the actor improves, the same starting position becomes genuinely more valuable, and the critic tracks it.

There is a ceiling. With GAMMA = 0.99 and a 500-step limit, the largest possible discounted return is about 99.3. The critic settles a little above that at 104.6, which is honest approximation error rather than a bug — the critic is a small network fitting noisy targets, and it overshoots slightly. If you saw 400 there, that would be a bug.

Notice also the dip at episode 400 (236.8) before the recovery. Actor-critic on CartPole is more stable than DQN and it is not monotone.

Line by line#

advantage = g - values — the whole method. g is what actually happened. values is what the critic expected. The difference is the advantage.

advantage.detach() — this is the line people get wrong, and it is silent when wrong. Without detach, the actor's loss sends gradients into the critic, and the optimiser discovers it can reduce the actor loss by making the critic's predictions worse. The critic is corrupted, the advantage becomes meaningless, and training limps along without any error message. Detach the advantage before it touches the actor.

One optimiser over both parameter lists — convenient here. In larger setups the critic usually wants a higher learning rate than the actor, so two optimisers is the more common arrangement.

0.5 * critic_loss — the critic's loss is on the scale of returns (up to about 100) while the actor's is on the scale of log-probabilities (around 1). Without a weight, the critic's gradients drown the actor's. That 0.5 is a real tuning knob, not decoration.

- 0.01 * entropy — the entropy bonus. Entropy measures how undecided the policy is. Subtracting it from the loss rewards staying undecided. Without it, a policy that gets an early lead can collapse to always-push-left, and a deterministic policy has nothing left to explore with. Set the coefficient to 0.0 and you will see this happen on some seeds.

nn.functional.mse_loss(values, g) — the critic is doing plain supervised regression. Its inputs are states, its targets are the observed returns. Nothing about it is special; the difficulty is that its targets keep changing as the actor improves.

The version used in practice#

This script uses full-episode returns for the critic's target, which makes it a Monte-Carlo actor-critic. Real implementations bootstrap instead:

# one-step target: reward now, plus the critic's opinion of where we landed
target = reward + GAMMA * critic(next_state) * (1 - terminated)

That reduces variance sharply and adds bias. Generalised advantage estimation (GAE) blends the two with a parameter lam: lam = 1.0 gives the Monte-Carlo version above, lam = 0.0 gives the one-step version, and lam = 0.95 is the near-universal default. You will meet GAE fully in the PPO lesson.

Common mistakes#

Forgetting .detach() on the advantage. Described above. This is the most common actor-critic bug and it does not raise anything.

Sharing a network body between actor and critic without balancing the losses. A shared trunk saves compute and couples the two objectives. If the critic's gradients dominate, the shared features become good for value prediction and useless for choosing actions.

No entropy bonus. Premature collapse to a deterministic policy. Symptom: the score climbs, then freezes exactly, and the policy's entropy is near zero.

Using an advantage of zero mean but wild scale. Normalising advantages per batch ((adv - adv.mean()) / (adv.std() + 1e-8)) is standard and makes the actor's learning rate far less sensitive.

Bootstrapping through a time-limit truncation. When CartPole stops at 500 steps the pole is still up. Treating that as a terminal state teaches the critic that surviving 500 steps is worth nothing after step 500.

Try it yourself#

Set the entropy coefficient to 0.0 and run five seeds. On at least one, the policy will collapse early and the score will freeze low. Then set it to 0.1 and watch the opposite: the agent stays deliberately random and never commits, so the score plateaus in the middle. The right value is a balance, and finding it by hand is the fastest way to understand what entropy is buying you.

What to learn next#

• PPO — the actor-critic method that fixed the instability well enough to be a default.
• Policy gradients — the estimator the critic is reducing the variance of.
• RLHF — actor-critic applied to language models.