Actor-critic methods and generalized advantage estimation

What it is and when you reach for it

An actor-critic method is a policy gradient method that carries a second network. The actor is the policy π_θ(a|s), the thing you actually want; the critic is a value function V_φ(s) whose only job is to reduce the variance of the actor's gradient. On the policy gradients page the score-function estimator weights each action's log-probability gradient by the return that followed it, and the whole story there is variance: the return of a trajectory mixes the consequences of every action in it, so the signal attached to any one action is mostly noise contributed by the others. The critic attacks this directly. If you know the value of a state, you can ask a sharper question about each action: not "was this trajectory good" but "was this action better than what I expected from this state". That difference is the advantage, and estimating it well is the entire subject of this page. You reach for an actor-critic whenever REINFORCE-style updates are too noisy to make progress in reasonable time, which in practice means almost always: nearly every modern policy optimization method, PPO included, is an actor-critic under the hood, and even critic-free methods like the one on the GRPO page are best understood as answers to the question "what if we replace the learned critic with something cheaper".

The math

The actor-critic decomposition

The policy gradient theorem says the gradient of expected return J(θ) can be written

∇_θJ(θ) = E[ Σ_t ∇_θ log π_θ(a_t|s_t) · Ψ_t ],

where Ψ_t is any quantity whose conditional expectation given (s_t, a_t) equals the action's true value contribution, possibly shifted by a baseline that depends only on the state. The choices for Ψ_t form a ladder. The full return G_t = Σ_k≥0 γ^kr_t+k gives REINFORCE: unbiased, maximally noisy. Subtracting a state baseline b(s_t) leaves the expectation untouched, because E_a∼π[∇ log π(a|s)] = 0 makes any action-independent term vanish in expectation, and the variance-minimizing baseline is approximately the state value V^π(s_t). One rung further, replace the noisy return entirely with learned estimates: the state-action value Q^π(s_t, a_t), or better, the advantage

A^π(s, a) = Q^π(s, a) − V^π(s),

the amount by which taking action a and then following the policy beats the policy's average from s. The advantage is the natural currency of policy improvement: it is zero in expectation over the policy's own actions, positive exactly for the actions worth reinforcing, and it yields the lowest-variance unbiased weighting in this family. The actor-critic decomposition is the decision to estimate V with a learned critic and to move the actor in the direction of the estimated advantage: the critic answers "what did I expect here", the actor updates on "how much did this action beat that expectation". Nothing about the actor's update rule changes from vanilla policy gradient; only Ψ_t does.

The TD error is a one-sample advantage estimate

The advantage involves Q, and learning a separate Q network would double the machinery. The temporal-difference error gets it from V alone. Define

δ_t = r_t + γV(s_t+1) − V(s_t).

Take its expectation over the environment's transition given (s_t, a_t), assuming V is the true value function V^π: the first two terms average to E[r_t + γV^π(s_t+1)] = Q^π(s_t, a_t) by the Bellman equation, so

E[δ_t | s_t, a_t] = Q^π(s_t, a_t) − V^π(s_t) = A^π(s_t, a_t).

One reward plus one value lookup is an unbiased sample of the advantage, provided the critic is exact. It never is, and that is the catch: with an approximate V, δ_t is a biased estimate whose bias is inherited from the critic's error at two states. In exchange the variance is tiny, since only a single reward's worth of environment noise enters. The TD error is therefore one extreme of a spectrum: maximum bias, minimum variance.

n-step returns walk the spectrum

The opposite extreme is the Monte Carlo advantage estimate G_t − V(s_t): no reliance on the critic beyond the baseline, so no bias from critic error, but every future reward's noise included. Between the two sit the n-step advantage estimates,

Â_t⁽ⁿ⁾ = r_t + γr_t+1 + … + γⁿ⁻¹r_t+n−1 + γⁿV(s_t+n) − V(s_t),

which trust real rewards for n steps and then hand off to the critic. n = 1 recovers the TD error; n → ∞ recovers Monte Carlo. Larger n means less bias, because the critic's error is discounted by γⁿ before it enters, and more variance, because n steps of reward and transition noise accumulate. A useful identity makes the whole family compact: telescoping the intermediate values gives

Â_t⁽ⁿ⁾ = Σ_l=0ⁿ⁻¹ γ^l δ_t+l.

Every n-step advantage is a discounted sum of one-step TD errors. Verifying for n = 2: δ_t + γδ_t+1 = (r_t + γV(s_t+1) − V(s_t)) + γ(r_t+1 + γV(s_t+2) − V(s_t+1)) = r_t + γr_t+1 + γ²V(s_t+2) − V(s_t), with the intermediate V(s_t+1) cancelling. This identity is what makes the next step natural.

Generalized advantage estimation

Rather than commit to one n, GAE (Schulman et al., 2015) takes an exponentially weighted average of all of them, with weight (1 − λ)λⁿ⁻¹ on the n-step estimate:

Â_t^GAE(γ,λ) = (1 − λ) Σ_n≥1 λⁿ⁻¹ Â_t⁽ⁿ⁾.

Substitute the telescoped form and swap the order of summation. The TD error δ_t+l appears in every n-step estimate with n > l, so its total weight is γ^l(1 − λ)(λ^l + λ^l+1 + …) = γ^lλ^l, and the whole average collapses to a single geometric sum:

Â_t^GAE(γ,λ) = Σ_l≥0 (γλ)^l δ_t+l.

λ is a bias-variance dial with the two familiar endpoints: λ = 0 gives exactly the TD error (biased, low variance), λ = 1 gives exactly the Monte Carlo advantage G_t − V(s_t) (unbiased given the baseline, high variance), and values in between, typically 0.9 to 0.98, buy most of the variance reduction at a modest bias cost. The estimator also admits a backward recursion, which is how everyone computes it: from the geometric structure,

Â_t = δ_t + γλ Â_t+1,

initialized with Â_T past the end of the rollout equal to zero, and with the recursion cut (multiplied by zero) wherever an episode terminated, since no reward flows across a reset. The value targets for training the critic come for free: R_t = Â_t + V(s_t) is exactly the λ-return, so one backward pass produces both the actor's weights and the critic's regression targets.

A worked example, small enough to check by hand. Take a three-step rollout with all rewards 1, γ = 0.9, λ = 0.8, critic outputs V(s₀) = 4.0, V(s₁) = 3.5, V(s₂) = 3.0, and a bootstrap value V(s₃) = 2.5 for the state after the rollout ends (no terminations). The TD errors, computed forward: δ₀ = 1 + 0.9·3.5 − 4.0 = 0.15, δ₁ = 1 + 0.9·3.0 − 3.5 = 0.20, δ₂ = 1 + 0.9·2.5 − 3.0 = 0.25. The advantages, computed backward with γλ = 0.72: Â₂ = 0.25, Â₁ = 0.20 + 0.72·0.25 = 0.38, Â₀ = 0.15 + 0.72·0.38 = 0.4236. Every advantage is positive, meaning the critic is currently underestimating these states, and the actor will be nudged toward all three actions, most strongly the earliest one, whose estimate pools the most evidence. The value targets are 4.4236, 3.88, and 3.25, each above the critic's current output, so the critic will be pulled up toward consistency at the same time.

A2C: the synchronous update

A3C (Mnih et al., 2016) originally ran many actor threads updating a shared network asynchronously; A2C is the later observation, popularized by OpenAI, that the asynchrony was an engineering workaround rather than an algorithmic ingredient, and that stepping the environments in lockstep and applying one synchronous batched update is simpler, more reproducible, and better suited to GPUs. The A2C update is: collect a short rollout (a few dozen steps) with the current policy, compute GAE advantages and value targets, and take one gradient step on the combined loss

L(θ, φ) = −E[ log π_θ(a_t|s_t) Â_t ] + c_v E[ (V_φ(s_t) − R_t)² ] − c_e E[ H(π_θ(·|s_t)) ],

with the advantage treated as a constant with respect to θ and φ (it is computed from stored values and never carries gradient), a value coefficient c_v around 0.5, and a small entropy bonus c_e around 0.01 that resists premature collapse onto a deterministic policy while there is still exploring to do. One rollout, one gradient step, throw the data away: A2C is strictly on-policy, and its data inefficiency is precisely the itch that PPO scratches by squeezing several epochs out of each batch.

Implementation, twice

Both implementations train A2C with GAE on CartPole-v1 using gymnasium, with a shared trunk feeding two heads: policy logits for the actor and a scalar value for the critic. Sharing the trunk halves the parameters and lets both objectives shape the same features, which helps on small problems; the value loss coefficient keeps the critic's regression gradients from drowning the policy gradient in the shared layers. The GAE computation is where the two frameworks diverge most instructively. In PyTorch it is a plain Python loop running backward over the rollout, which is fine because the rollout is short and the loop is outside autograd. In JAX the same recursion is a lax.scan with reverse=True, and this is worth pausing on: GAE is a linear recurrence scanned backward over a trajectory, which is exactly the shape of computation lax.scan exists for, so the JAX version is not a workaround but the estimator's most natural expression: the recursion Â_t = δ_t + γλÂ_t+1 written as a one-line scan body, compiled once, differentiable if you ever need it to be.

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

GAMMA, LAM = 0.99, 0.95
N_STEPS = 32                  # short rollout per update, A2C style
LR, ENT_COEF, VF_COEF = 7e-4, 0.01, 0.5
TOTAL_STEPS = 200_000

class ActorCritic(nn.Module):
    """Shared trunk, two heads: logits for the actor, scalar V for the critic."""
    def __init__(self, obs_dim, n_actions, hidden=128):
        super().__init__()
        self.trunk = nn.Sequential(
            nn.Linear(obs_dim, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh())
        self.pi = nn.Linear(hidden, n_actions)
        self.v = nn.Linear(hidden, 1)

    def forward(self, x):
        h = self.trunk(x)
        return self.pi(h), self.v(h).squeeze(-1)

def gae(rewards, values, dones, last_value):
    """A_t = delta_t + gamma*lam*(1 - done_t)*A_{t+1}, walked backward.
    done_t cuts the recursion so no credit flows across an episode reset."""
    adv = torch.zeros_like(rewards)
    a, next_v = 0.0, last_value
    for t in reversed(range(len(rewards))):
        mask = 1.0 - dones[t]
        delta = rewards[t] + GAMMA * next_v * mask - values[t]
        a = delta + GAMMA * LAM * mask * a
        adv[t] = a
        next_v = values[t]
    return adv, adv + values          # advantages, critic targets (lambda-returns)

env = gym.make("CartPole-v1")
net = ActorCritic(env.observation_space.shape[0], env.action_space.n)
opt = torch.optim.Adam(net.parameters(), lr=LR)

obs, _ = env.reset(seed=0)
ep_ret, recent = 0.0, []
for update in range(TOTAL_STEPS // N_STEPS):
    obs_b, act_b, rew_b, done_b, val_b = [], [], [], [], []
    for _ in range(N_STEPS):
        with torch.no_grad():         # rollout carries no gradient
            logits, value = net(torch.as_tensor(obs, dtype=torch.float32))
        action = torch.distributions.Categorical(logits=logits).sample()
        next_obs, r, terminated, truncated, _ = env.step(action.item())
        obs_b.append(obs); act_b.append(action)
        rew_b.append(r); val_b.append(value)
        # Folding truncation into done is the common simplification; it
        # slightly biases values at the 500-step time limit (see traps).
        done_b.append(float(terminated or truncated))
        ep_ret += r
        obs = next_obs
        if terminated or truncated:
            recent.append(ep_ret); ep_ret = 0.0
            obs, _ = env.reset()
    with torch.no_grad():             # bootstrap value for the cut-off tail
        _, last_value = net(torch.as_tensor(obs, dtype=torch.float32))

    obs_t = torch.as_tensor(np.array(obs_b), dtype=torch.float32)
    act_t = torch.stack(act_b)
    adv, ret = gae(torch.as_tensor(rew_b, dtype=torch.float32),
                   torch.stack(val_b),
                   torch.as_tensor(done_b), last_value)

    logits, value = net(obs_t)        # recompute with gradients on
    dist = torch.distributions.Categorical(logits=logits)
    pg_loss = -(dist.log_prob(act_t) * adv).mean()   # adv is a constant here
    v_loss = F.mse_loss(value, ret)
    entropy = dist.entropy().mean()
    loss = pg_loss + VF_COEF * v_loss - ENT_COEF * entropy

    opt.zero_grad()
    loss.backward()
    nn.utils.clip_grad_norm_(net.parameters(), 0.5)
    opt.step()

    if update % 200 == 0 and recent:
        print(f"step {update * N_STEPS:>7}  "
              f"mean return {np.mean(recent[-20:]):.1f}")

import gymnasium as gym
import numpy as np
import jax
import jax.numpy as jnp
import optax

GAMMA, LAM = 0.99, 0.95
N_STEPS = 32                  # short rollout per update, A2C style
LR, ENT_COEF, VF_COEF = 7e-4, 0.01, 0.5
TOTAL_STEPS = 200_000

def dense(key, n_in, n_out):
    scale = jnp.sqrt(2.0 / n_in)
    return {'w': jax.random.normal(key, (n_in, n_out)) * scale,
            'b': jnp.zeros(n_out)}

def init_params(key, obs_dim, n_actions, hidden=128):
    k1, k2, k3, k4 = jax.random.split(key, 4)
    return {'l1': dense(k1, obs_dim, hidden), 'l2': dense(k2, hidden, hidden),
            'pi': dense(k3, hidden, n_actions), 'v': dense(k4, hidden, 1)}

def forward(params, x):
    """Shared trunk, two heads: logits for the actor, scalar V for the critic."""
    h = jnp.tanh(x @ params['l1']['w'] + params['l1']['b'])
    h = jnp.tanh(h @ params['l2']['w'] + params['l2']['b'])
    logits = h @ params['pi']['w'] + params['pi']['b']
    value = (h @ params['v']['w'] + params['v']['b']).squeeze(-1)
    return logits, value

def gae(rewards, values, dones, last_value):
    """GAE is a linear recurrence walked backward over the trajectory,
    which is exactly what lax.scan(reverse=True) is for: the scan body
    IS the recursion A_t = delta_t + gamma*lam*(1-done)*A_{t+1}, and
    outputs come back already in forward time order."""
    def step(carry, x):
        a, next_v = carry
        r, v, d = x
        mask = 1.0 - d
        delta = r + GAMMA * next_v * mask - v
        a = delta + GAMMA * LAM * mask * a
        return (a, v), a
    _, adv = jax.lax.scan(step, (jnp.float32(0.0), last_value),
                          (rewards, values, dones), reverse=True)
    return adv, adv + values          # advantages, critic targets

optimizer = optax.chain(optax.clip_by_global_norm(0.5), optax.adam(LR))

def loss_fn(params, obs, actions, adv, returns):
    logits, values = forward(params, obs)
    logp = jax.nn.log_softmax(logits)
    chosen = jnp.take_along_axis(logp, actions[:, None], axis=1).squeeze(1)
    pg_loss = -(chosen * adv).mean()  # adv enters as data, not a function of params
    v_loss = jnp.mean((values - returns) ** 2)
    entropy = -jnp.sum(jnp.exp(logp) * logp, axis=1).mean()
    return pg_loss + VF_COEF * v_loss - ENT_COEF * entropy

@jax.jit
def update(params, opt_state, obs, actions, rewards, dones, values, last_value):
    adv, returns = gae(rewards, values, dones, last_value)
    grads = jax.grad(loss_fn)(params, obs, actions, adv, returns)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state

@jax.jit
def act(params, obs, key):
    logits, value = forward(params, obs)
    return jax.random.categorical(key, logits), value

env = gym.make("CartPole-v1")
key = jax.random.PRNGKey(0)
params = init_params(key, env.observation_space.shape[0], env.action_space.n)
opt_state = optimizer.init(params)

obs, _ = env.reset(seed=0)
ep_ret, recent = 0.0, []
for upd in range(TOTAL_STEPS // N_STEPS):
    obs_b, act_b, rew_b, done_b, val_b = [], [], [], [], []
    for _ in range(N_STEPS):
        key, k = jax.random.split(key)
        action, value = act(params, jnp.asarray(obs), k)
        next_obs, r, terminated, truncated, _ = env.step(int(action))
        obs_b.append(obs); act_b.append(int(action))
        rew_b.append(r); val_b.append(float(value))
        done_b.append(float(terminated or truncated))
        ep_ret += r
        obs = next_obs
        if terminated or truncated:
            recent.append(ep_ret); ep_ret = 0.0
            obs, _ = env.reset()
    _, last_value = forward(params, jnp.asarray(obs))

    params, opt_state = update(
        params, opt_state,
        jnp.asarray(np.array(obs_b), dtype=jnp.float32),
        jnp.asarray(act_b), jnp.asarray(rew_b, dtype=jnp.float32),
        jnp.asarray(done_b, dtype=jnp.float32),
        jnp.asarray(val_b, dtype=jnp.float32), last_value)

    if upd % 200 == 0 and recent:
        print(f"step {upd * N_STEPS:>7}  "
              f"mean return {np.mean(recent[-20:]):.1f}")

A few deliberate choices. The rollout is collected without gradients and the network is re-run on the batch for the update; for a 32-step batch the recomputation is free, and it keeps the collection loop clean of autograd bookkeeping. The advantages are computed from stored values and never carry gradient into the policy loss, which is not an optimization detail but a correctness requirement (see traps). Gradient clipping at global norm 0.5 and the tanh trunk both come straight from what the reference libraries do for this class of method. In the JAX version everything that touches arrays lives in jitted functions, and the environment stepping stays in Python where it must, since gymnasium is stateful, host-side code; the split between "pure, jitted math" and "impure, sequential world" is the standard architecture for JAX RL against non-JAX environments.

Using it on a real shape of problem

CartPole-v1 has a 4-dimensional observation, 2 actions, reward 1 per step, and a 500-step time limit, so a solved policy earns a return near 500. Run either implementation as-is. The behavior to expect: mean return hovers around 20 to 40 for the first few thousand steps while the critic is still learning what "expected" means, then climbs, often noisily and with occasional collapses of 50 to 100 return that recover within a few thousand steps, and reaches the 300 to 500 range somewhere in the second half of the 200k-step budget; in my own runs of these exact files one seed crossed 400 near 115k steps and another was still around 330 at 150k. The exact trajectory is seed- and machine-dependent; single-environment A2C on CartPole has real run-to-run variance, and an unlucky seed can take noticeably longer. Two diagnostics are worth watching beyond the return. The value loss should fall over training but never to zero, since the targets keep moving as the policy improves. The entropy should decay smoothly from ln 2 ≈ 0.693 toward zero as the policy commits; if it crashes to zero in the first few thousand steps the policy has collapsed before finding anything worth committing to, and the usual fix is a larger entropy coefficient or a smaller learning rate.

Applications

The honest framing is that A2C itself is rarely the final answer anymore, but the actor-critic-with-GAE template it embodies is the chassis of modern policy optimization. PPO is literally this page's algorithm with a clipped objective and multiple epochs per batch; everything else, the two-headed network, the rollout buffer, GAE, the entropy bonus, transfers unchanged, which is why reading this page first makes the PPO page short. The same template, with PPO's objective, trained OpenAI Five for Dota 2 and the dexterous robot hand work, drives most of the locomotion and manipulation results in simulated robotics (thousands of parallel environments stepping in lockstep is A2C's synchronous idea taken to its extreme on GPU simulators like Isaac Gym), and underpins game-playing agents from Atari benchmarks onward. In RLHF for language models the actor is the LLM itself, the critic is a value head grafted onto it, and GAE runs over token sequences (generating those rollouts at scale is a systems problem of its own, covered in my vLLM notes); and when GRPO drops the critic to halve the memory bill, the thing it replaces, a learned baseline for advantage estimation, is exactly the machinery derived above, which is why GRPO's group-mean baseline is best judged against GAE's bias-variance analysis. Asynchronous descendants of A3C, notably IMPALA with its V-trace correction, run large-scale distributed training at places where a single synchronous batch cannot keep thousands of actors busy. If you learn one on-policy algorithm's anatomy deeply, it should be this one, because every neighbor is a small edit of it.

Against the real libraries

stable-baselines3 ships A2C as a first-class algorithm, and its implementation is this page's plus the production trimmings: vectorized environments stepping n copies in lockstep so each update sees a (n_envs × n_steps) batch, RMSprop with the specific epsilon and alpha the original A3C paper tuned (a detail that measurably matters for reproducing published curves; SB3 defaults to RMSprop where this page uses Adam), optional observation and reward normalization via VecNormalize, correct bootstrapping at time-limit truncations, learning-rate schedules, and logging. Its RolloutBuffer.compute_returns_and_advantage is the same backward recursion as the gae functions above, which makes it a good verification target: collect a fixed rollout, feed identical rewards, values, and dones to both implementations, and the advantages should match to float precision. The concrete check: seed everything, fill an SB3 RolloutBuffer with your arrays, call its compute_returns_and_advantage, and assert np.allclose(buffer.advantages.flatten(), yours, atol=1e-6); when it disagrees, the culprit is done-masking or bootstrap handling essentially every time.

CleanRL is the other reference worth reading, for the opposite reason: instead of a framework it gives you single-file implementations (its ppo.py line, plus PyTorch and JAX variants of the classic algorithms) where every detail of this template, GAE, advantage handling, the two-head network, is visible in one scroll. CleanRL's JAX implementations use exactly the lax.scan GAE pattern shown above, and its documentation of seeded, benchmarked runs gives you honest learning curves to compare against rather than a vague sense that "it should work". When is the from-scratch version enough? For single-environment problems of CartPole's scale, for teaching, and for research where you intend to modify the estimator itself (a new λ schedule, a different baseline), where a framework's abstraction is friction. The moment you need parallel environments, checkpointing, normalization, and evaluation protocols, the libraries stop being convenience and start being correctness, because each of those features hides a handful of the traps below.

Traps and misconceptions

Letting gradient flow through the advantage. The policy loss is −log π(a|s) ·  with  as data. If  is computed from value tensors that still carry graph, the policy update differentiates through the critic and the actor starts optimizing the critic's opinion of it rather than acting on it, which produces confidently wrong learning rather than an error message. Compute values under no_grad (or stop the gradient explicitly) before they enter the GAE recursion; in the JAX version this is structural, since advantages enter loss_fn as arguments that are never functions of params.

Treating time-limit truncation as termination. CartPole ends at 500 steps by fiat, not because the state was bad. Setting done = 1 there tells GAE the future was worth zero, which systematically underestimates values near the limit. The implementations above accept this small bias for simplicity and say so; the correct treatment, which stable-baselines3 implements, bootstraps γV(s_t+1) at truncations while cutting only at true terminations. gymnasium returns terminated and truncated separately precisely so you can make this distinction; code that merges them without thinking inherits the bias on every time-limited environment.

"λ = 1 is best because it is unbiased." Unbiased is not the objective; low mean-squared error of the gradient is. At λ = 1 the estimator ignores the critic's variance reduction entirely and you have paid for a critic you are not using. The empirically strong region, λ around 0.9 to 0.98, accepts bias the critic's accuracy can cover in exchange for a large variance cut. Symmetrically, λ = 0 leans entirely on a critic that is wrong for most of training. The dial exists because both endpoints are bad.

Ignoring the shared-trunk gradient balance. With one trunk serving both heads, the value regression can produce gradients much larger than the policy term, especially early, when targets are far from predictions. That is what the value coefficient (0.5 here) and gradient clipping are managing. If the policy seems inert while value loss falls, the critic is monopolizing the trunk; lower c_v, clip harder, or split the networks, which costs parameters but removes the interference entirely.

Expecting the critic to converge first. The critic is not a prerequisite the way a target network is in DQN; actor and critic improve together, and the critic chases a moving target defined by the current policy forever. A critic that fits its targets perfectly mid-training is a warning sign (usually value overfitting on a small buffer), not a milestone. The system works because the advantage only needs to rank actions better than noise would, long before V is accurate in absolute terms.