What it actually does
REINFORCE weights each action by the return that actually followed it, which is unbiased and enormously noisy, because a single sampled return over a long horizon is dominated by everything that happened to occur afterwards. The fix is already visible in the theory: any function of state can be subtracted for free, and the best one is the state's own value. Learn that function and two things happen at once. The subtraction removes the part of the return that was predictable from the state, leaving only the part attributable to the action — which is the definition of an advantage. And once a value function exists, the sampled return can be replaced entirely by a bootstrapped estimate, so an update no longer has to wait for the episode to end. What you gain is a dial. Bootstrap after one step and variance nearly vanishes while bias is whatever the critic's error is; bootstrap after the whole episode and you are back to REINFORCE. Generalized advantage estimation makes that dial continuous, and choosing where to set it is the central tuning decision of the method. What you pay is a second learning problem coupled to the first. The critic supplies the actor's training signal, the actor determines the distribution the critic is fitted on, and neither is right at the start. An actor that improves faster than its critic is optimizing against systematically wrong advantages, which is the characteristic failure and the reason relative learning rates matter more here than anywhere else in this category.
- The policy is differentiable and assigns non-zero probability to every action it might need, since an action never sampled is never improved
- The critic can represent the value function well enough on the states the policy visits — everywhere else its advantages are extrapolation
- Trajectories come from the current policy, which is what makes the advantage estimate valid and what forbids replay
- Environments can be replicated in parallel, because that is where the decorrelation comes from in the absence of a replay buffer
- The critic adapts at least as fast as the actor, which the two-timescale results require and which the learning-rate ratio must actually deliver
- Rewards are on a bounded or normalized scale, since their magnitude sets the scale of both the advantage and the value regression
Loss function
Three objectives are being optimized in one expression, and they are not the same kind of object. The first term is a policy-gradient surrogate whose value is meaningless and whose gradient is the point. The second is an ordinary regression with an interpretable loss — the critic is genuinely being fitted to a target and its error genuinely measures something. The third is a regularizer keeping the policy from committing prematurely. Reading them as one number is a mistake: a total loss that falls can be a critic improving while the policy degrades, and the two must be logged separately or the run cannot be diagnosed. The advantage estimate is where the real design decision lives. A one-step advantage r + gamma·V(s') − V(s) has almost no variance and carries all of the critic's bias; the full Monte Carlo advantage has no bias beyond the baseline and all of REINFORCE's variance. Generalized advantage estimation interpolates between them with an exponentially weighted sum of TD errors, and its lambda is a bias-variance dial that is usually worth more attention than the learning rate. The detachment requirement deserves to be stated as loudly as possible: the advantage appears in the policy term as a scalar weight, and if it is left attached to the computation graph, the policy loss backpropagates into the critic and trains it to make advantages large rather than accurate. Nothing errors, the loss curves look plausible, and the run is silently optimizing the wrong thing.
- the advantage estimate, treated as a constant by the policy term — the gradient must never flow through it into the critic, and this is the single most common bug in the method
- the critic: a learned state-value function, fitted by regression onto bootstrapped returns, whose only job is to make the advantage less noisy
- the value-loss coefficient, which matters mostly because a shared trunk lets the regression gradient overwhelm the far smaller policy gradient
- the entropy coefficient; the policy is still the sole source of exploration, so collapse is still permanent
How it is minimized
Simultaneous stochastic gradient descent on a coupled policy surrogate and value regression, with advantages estimated by GAE over short on-policy rollouts
- the one-step TD error — the same quantity TD learning updates on, reused here as the unit the advantage is built from
- the bias-variance dial: 0 gives the low-variance, fully bootstrapped advantage, 1 gives the unbiased Monte Carlo one
- the effective decay on how far forward TD errors are accumulated, which sets the horizon the advantage actually sees
- the bootstrap; it must be zeroed at a true terminal and retained at a time-limit truncation, and conflating those is a standard bug
The method exists to trade REINFORCE's variance for a controlled amount of bias, and almost every design decision is about managing the coupling that trade introduces. GAE is the mechanism: accumulating discounted TD errors over the rollout gives an advantage whose bias and variance are set by a single parameter, computed in one reversed pass for free. The relative learning rates of actor and critic are the next decision and the one most often got wrong — an actor moving faster than its critic optimizes against advantages that are systematically wrong, and the usual symptom is a return curve that rises and then falls away with no accompanying change in the loss. A shared trunk between actor and critic is common because features useful for valuing a state are useful for choosing in it, and it introduces its own problem: the value regression has a much larger natural gradient scale than the policy surrogate, so without a value-loss coefficient the trunk is trained almost entirely by the critic. Rollout length is the third decision, and it is a genuine trade rather than a tuning detail — short rollouts update more often on staler bootstraps, long ones are closer to Monte Carlo and update less. A2C and A3C differ only in how parallelism is arranged: A3C ran asynchronous workers applying lock-free updates from slightly stale parameters, A2C steps many environments in lockstep and applies one synchronous update, and the synchronous version turned out to be both simpler and at least as good, which is a useful lesson about where the benefit of the asynchronous design actually came from. It was the decorrelation from parallel environments, not the asynchrony.
GAE lambda
The bias-variance dial on the advantage. Often worth more than the learning rate, and the first thing to move when the gradient is too noisy or the critic is too trusted (0.9 to 0.97)
value-loss coefficient
Balances a regression with a large natural gradient scale against a policy surrogate with a small one. Matters most with a shared trunk, where the critic otherwise trains it alone (0.25 to 1.0)
entropy coefficient
The policy is still the only source of exploration, so collapse is still permanent. Unchanged in role from REINFORCE and unchanged in importance (0.0 to 0.01)
rollout length
Steps collected per environment before an update. Short means frequent updates on staler bootstraps; long approaches Monte Carlo and updates less often (5 to 128 steps)
number of parallel environments
The decorrelation mechanism, and the actual source of A3C's benefit. Independent environments make a batch approximately i.i.d. without any replay buffer (8 to 64)
actor and critic learning rates
Their ratio is the coupling knob. A critic that lags leaves the actor optimizing against wrong advantages, which is the characteristic failure of the method (3e-4 actor, equal or up to 3x for the critic)
discount (gamma)
The horizon, and jointly with lambda it sets how far forward the advantage actually looks — the two interact and should not be tuned independently (0.99)
No convergence guarantee worth relying on. Actor-critic is a two-timescale stochastic approximation, and results exist for the tabular and linear cases given a critic that adapts faster than the actor, but they do not survive nonlinear function approximation and the practical behaviour is governed by the coupling rather than by the theory. The characteristic failure follows directly from that coupling: when the actor improves faster than the critic can track, the advantages it optimizes against are systematically wrong, and the run produces a return curve that climbs, turns over, and collapses without the loss showing anything unusual. The diagnosis is the critic's explained variance against realized returns, which is the single most informative number in an actor-critic run and is routinely not logged. Below roughly zero the critic is worse than predicting the mean and the advantages are noise. Entropy collapse is inherited unchanged from REINFORCE and remains permanent — the critic does nothing to restore exploration, and a policy that has stopped sampling an action still never learns about it. A third failure is specific to the shared trunk: because the value regression has a much larger gradient scale, a badly set value coefficient produces a trunk that serves the critic and starves the policy, which presents as an actor that will not move. And as everywhere in this category, seed variance is large enough that a single-seed comparison measures the seed rather than the method.
Per environment step: one forward pass through the policy and one through the critic, or one shared forward with two heads. Per update: one backward over a rollout of T steps across N environments, so the batch is T·N transitions and there is no replay buffer to sample from. GAE costs a single reversed pass over the rollout, O(T·N), which is negligible. Memory is bounded by the rollout rather than by stored history, which makes this substantially lighter than DQN and is a real practical advantage on constrained hardware. Sample complexity sits between REINFORCE and an off-policy method: bootstrapping recovers much of the efficiency lost to being on-policy, but every batch is still discarded after one gradient step, which is precisely the waste PPO removes by taking several epochs over the same data.
How, where, and why it is used
Time-Series Forecasting
Not applicableThere is no action whose advantage could be estimated and no policy for a critic to evaluate, since a forecaster does not influence what it observes next; the critic here is a value function over a controlled process, not a predictor of a series.
Anomaly Detection
Not applicableDetection scores observations without acting, so the advantage — how much better an action turned out than expected — has nothing to attach to, and a large critic error reports an under-fitted region of the value function rather than anything unusual in the data.
Optimization
Primary fitOptimize a parameterized policy using advantages estimated by a learned critic rather than by sampled returns, trading a controlled amount of bias for a large reduction in variance and removing the need to wait for episodes to end.
- The bias-variance trade made explicit and tunable, with GAE's lambda as a single continuous dial between full bootstrapping and Monte Carlo
- Two coupled learners as an engineering problem: each supplies the other's signal, and their relative rates decide whether the run works
- Parallel environments as the decorrelation mechanism, which is what A3C's benefit actually came from and what A2C kept while dropping the asynchrony
- The detachment requirement — a scalar weight that must not carry gradient — as a concrete instance of a general hazard in multi-objective training
The synthesis point of this category, and the reason it is worth studying even though PPO has largely replaced it in deployment. Three lessons transfer past reinforcement learning. The first is that a learned baseline is the general way to reduce the variance of a sampled estimator, and that once a baseline exists it can usually also replace part of the sample — which is what bootstrapping does here and what control variates do elsewhere. The second is about coupled learners: whenever one model produces another's training signal, their relative learning rates become a first-class hyperparameter, and the failure when they are wrong is not an error but a plausible-looking run optimizing the wrong objective. The third is the detachment hazard — a quantity used as a weight rather than as a prediction must be severed from the graph, and forgetting it produces no error, no warning, and a silently different objective. Where actor-critic is the wrong choice: if the horizon is short and returns are cheap, REINFORCE with a batch baseline is simpler and nearly as good; if stability matters, PPO is the same method with a trust region and several epochs per batch, which dominates it; and for continuous control where sample efficiency is the binding constraint, an off-policy method such as SAC is a better use of the same interactions.
- Log the policy and value losses separately — summed, a rising policy loss hides inside a falling value loss
- Track the critic's explained variance against realized returns; below zero the advantages are noise and nothing downstream is meaningful
- Detach the advantage before it multiplies the log-probability, or the policy term trains the critic to make advantages large rather than correct
- Zero the bootstrap at a true terminal and keep it at a time-limit truncation, which are different events that a single done flag conflates
Median return over at least five seeds with the spread shown, evaluated on a separate deterministic rollout rather than on the exploring training policy. Plot critic explained variance and policy entropy alongside return: a plateau is almost always explained by one of the two — a critic that stopped tracking, or a policy that stopped exploring — and neither is visible in the return curve itself.
- A total loss reported as one number, which hides a degrading policy inside an improving critic
- An advantage left attached to the graph, which silently changes what the critic is trained to do
- An actor tuned to move faster than its critic, which optimizes against advantages that are systematically wrong
- Entropy collapse, inherited unchanged from REINFORCE and equally permanent — the critic does nothing to restore exploration
Control & Operations
ViableLearn a control policy over continuous or discrete actions from parallel simulated interaction, with a critic supplying low-variance advantages so that updates can happen every few steps instead of once per episode.
- Continuous-action control where a value-based argmax is intractable and Monte Carlo returns are too noisy to use directly
- Long-horizon operations problems where waiting for an episode to end would make updates impossibly infrequent
- Simulators that can be replicated cheaply, since parallel environments are the decorrelation mechanism this method depends on
- Settings where memory is constrained, because a rollout buffer is orders of magnitude smaller than a replay buffer
A real fit, and the honest framing is that it is the foundation rather than the destination. Everything this method brings to an operations problem — continuous actions, bootstrapped updates on long horizons, modest memory — is retained by PPO, which adds a trust region and several epochs per batch and is more robust to the hyperparameters that matter. So the practical recommendation is to understand actor-critic and deploy PPO, or SAC where sample efficiency binds hardest. Two cautions apply to the family and deserve stating here. The critic is a model of value under the current policy, so it is only trustworthy in the region the current policy visits, and a policy that moves into an unfamiliar operating regime is being advised by a critic extrapolating with confidence and no data. And a stochastic policy samples its actions by construction, which is a liability in a physical system: the deployed controller should be the distribution's mode, evaluated separately, and anything genuinely unsafe needs an interlock the policy cannot trade against.
- Normalize observations with running statistics, since both heads are sensitive to input scale and operating regimes differ by orders of magnitude
- Bound continuous actions explicitly, because an unbounded Gaussian policy will eventually sample something the actuator cannot execute
- Replicate the simulator rather than lengthening rollouts; parallel environments decorrelate a batch, longer rollouts only add horizon
- Deploy the mode and evaluate it separately, since a deterministic policy is a different controller from the stochastic one that was trained
Simulated return against the incumbent controller on matched scenarios, with constraint violations counted separately rather than averaged in. Report critic explained variance per operating regime as well as overall — a critic that is accurate on average and wrong in one regime produces a policy that is confidently bad in exactly that regime.
- Trusting the critic outside the state distribution the current policy actually visits
- Deploying the stochastic policy, so a production controller samples its actions
- Choosing this over PPO for anything beyond a study exercise, when PPO keeps the advantages and removes the main instabilities
- A hard constraint expressed as a reward penalty, which an optimizer treats as a price rather than a limit
Natural Language
AdaptedOptimize a generation policy against a sequence-level or preference-based reward with a learned value head supplying per-token advantages, which is the structure preference-tuning pipelines are built from.
- Preference optimization from a reward model, where the standard pipeline is an actor-critic method with a KL penalty to the reference policy
- Per-token credit assignment on a reward that only arrives at the end of the sequence, which is exactly what a critic is for
- A value head attached to the language model itself, sharing the trunk that already encodes the context
- Understanding why direct preference methods were proposed — they remove the critic, and with it this method's coupling problem
Marked adapted because what is actually deployed is PPO rather than plain A2C, but the structure is this one and the difficulties are this one's. The attraction is real: a sequence-level reward arriving only at the end gives every token the same credit under a Monte Carlo estimator, and a critic is precisely the machinery that distributes it across the tokens that earned it. The costs are also this method's costs, amplified by scale. Two large models train simultaneously with coupled signals, memory holds a policy, a critic, a reference and often a reward model at once, and the critic must be fitted on a distribution the policy is actively moving. The critic's explained variance is the diagnostic here as everywhere, and it is why direct preference optimization was attractive enough to displace this pipeline in many settings — it removes the critic entirely and therefore removes the coupling, at the cost of the flexibility a learned value function provides. Where the critic earns its place is where the reward is dense or intermediate credit genuinely matters.
- Start from a supervised-tuned policy; neither the actor nor the critic can discover language from a reward signal
- Keep a KL penalty to the reference policy, since an unconstrained optimizer finds the reward model's defects rather than the behaviour it encodes
- Attach the value head to the shared trunk but watch the value-loss coefficient, because the regression gradient will otherwise dominate the trunk
- Normalize rewards within a batch, as sequence-level scores are unbounded and their scale directly sets gradient scale
Held-out preference win rate alongside KL from the reference policy: a reward gain accompanied by a large divergence is usually reward hacking rather than improvement. Track critic explained variance as well — a critic that is not tracking makes every per-token advantage noise, and the run degenerates toward a worse version of the Monte Carlo estimator it was meant to improve on.
- Optimizing a learned reward without a KL constraint, which yields text that scores well and reads badly
- A value-loss coefficient left at a default on a shared trunk, so the critic trains the representation and the policy barely moves
- Assuming the critic is tracking without measuring it, when its failure is invisible in the reward curve
- Carrying the full pipeline when a direct preference method would remove the critic and the coupling along with it
What it costs to run
Moderate by the standards of this category: one or two forward passes per environment step and one backward pass per rollout, with no replay buffer and memory bounded by the rollout rather than by stored history. That makes it markedly lighter than DQN on memory and heavier on interaction, since every batch is discarded after a single gradient step. Parallel environments are the main scaling lever and the main cost — simulator instances multiply with worker count. Budget for at least five seeds, because the coupling between actor and critic makes run-to-run variance large. Illustrative, not a measured benchmark.
One forward pass through the policy; the critic is training machinery and is discarded at deployment. The served artefact is therefore smaller than what was trained, which is a genuine operational advantage. The decision that matters is whether to ship the distribution's mode or a sample: the trained object is stochastic, most deployments should not be, and the deterministic policy is a different controller whose performance must be measured rather than assumed.
Batch retraining rather than continuous updating, since the estimator is on-policy and live updating means a production system taking sampled actions with a critic that may have stopped tracking. The usual arrangement is periodic offline training against a simulator, offline evaluation of the deterministic policy, then a staged rollout. As with every simulator-trained policy, the simulator goes stale before the policy does, so revalidating it is the real trigger for a retraining cycle.
- Critic explained variance against realized returns — the single most informative number in an actor-critic run, and the one most often not logged; below zero, every advantage downstream is noise
- Policy and value losses separately, never summed, since a degrading policy hides comfortably inside an improving critic
- Policy entropy, which is still the leading indicator of permanent exploration collapse and is unaffected by the presence of a critic
- KL divergence between successive policy versions, because a large step both collects worse data and outruns the critic that was fitted on the old distribution
- Realized return of the deployed deterministic policy against the incumbent, on matched scenarios rather than in aggregate
- State-distribution shift between training and production, since the critic is only trustworthy where the training policy actually visited
- The advantage must be detached before it multiplies the log-probability. Left attached, the policy term backpropagates into the critic and trains it to produce large advantages rather than accurate ones — no error, plausible curves, wrong objective
- A summed loss is undiagnosable. The policy surrogate's value means nothing and the value loss means something, so reporting their sum discards the only interpretable half
- If the actor outruns the critic, the whole run optimizes against systematically wrong advantages, and the symptom is a return curve that rises and then collapses with nothing unusual in the loss
- With a shared trunk the value regression has a much larger gradient scale than the policy surrogate, so a default value coefficient can leave the trunk trained almost entirely by the critic
- A true terminal and a time-limit truncation require different bootstrap handling, and a single done flag conflates them — the resulting policy is visibly short-sighted at the episode limit
- Entropy collapse is inherited from REINFORCE unchanged. Adding a critic reduces variance and does nothing whatsoever to restore lost exploration
- The critic is only valid on the distribution the current policy visits, so a policy that moves regions is being advised by confident extrapolation
- A2C and A3C differ only in synchronization, and the benefit attributed to asynchrony was really decorrelation from parallel environments — worth knowing before building the harder one
Pros and cons, in context
A tunable bias-variance trade rather than a fixed one
GAE's lambda moves continuously between the fully bootstrapped advantage and the Monte Carlo one, computed in a single reversed pass. It is usually the most valuable parameter in the method and it costs nothing to compute
Updates without waiting for episodes to end
Bootstrapping removes REINFORCE's requirement for a complete return, which makes long-horizon and continuing tasks tractable and lets updates happen every few steps instead of once per episode
Keeps the policy-gradient advantages while recovering much of the lost efficiency
Continuous and structured action spaces, non-differentiable environments, and direct optimization of the deployed policy all survive, with far less of the variance that made the plain estimator unusable
Memory bounded by the rollout, not by stored history
No replay buffer at all, so the memory footprint is orders of magnitude below a value-based method of comparable capability. On constrained hardware this is frequently the deciding factor
The critic is training machinery and does not ship
Only the policy is served, so the deployed artefact is smaller than what was trained — the opposite of a value-based method, where the value function is the policy
Two coupled learners, each supplying the other's signal
Their relative learning rates become a first-class hyperparameter, and getting them wrong produces a plausible-looking run optimizing systematically wrong advantages. This is the defining difficulty of the method
The bias the critic introduces is real and unmeasured by default
Bootstrapping through an inaccurate value function biases every advantage, and nothing surfaces it unless explained variance is tracked — which it usually is not
The detachment bug is silent and easy to write
An advantage left on the graph makes the policy loss train the critic toward large advantages rather than accurate ones. No error, no warning, and curves that look entirely normal
Still on-policy, so data is still discarded after one step
Bootstrapping recovers efficiency relative to REINFORCE but every batch is used once. Taking several epochs over the same rollout is precisely what PPO adds, and it is why PPO superseded this
Entropy collapse is inherited unchanged
The critic reduces variance and restores no exploration whatsoever. A policy that has stopped sampling an action still never learns about it, and the failure remains permanent
More moving parts than its successors, with less robustness
Value coefficient, entropy coefficient, lambda, rollout length, worker count and two learning rates, all interacting. PPO has most of the same knobs and tolerates bad settings far better, which is why plain A2C is now mostly a teaching artefact
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""Actor-critic - the one-step advantage update, transcribed.
Two learners in one loop. The critic is fitted by regression onto a
bootstrapped return; the actor is pushed by a policy gradient weighted by how
wrong the critic turned out to be.
Watch for the coincidence that makes the method this compact: the SAME
quantity serves both. The TD error is the critic's regression residual and the
actor's advantage estimate, computed once and used twice.
"""
import math
import random
def softmax(scores):
largest = max(scores)
exponentials = [math.exp(score - largest) for score in scores]
total = sum(exponentials)
return [value / total for value in exponentials]
def policy_probabilities(theta, state):
"""theta[a][i] is the actor's weight on feature i for action a."""
scores = []
for action_weights in theta:
score = 0.0
for weight, feature in zip(action_weights, state):
score += weight * feature
scores.append(score)
return softmax(scores)
def state_value(w, state):
"""A linear critic: V(s) = w . s.
Linear on purpose - its gradient with respect to w is just s, so the
critic's update stays readable next to the actor's.
"""
total = 0.0
for weight, feature in zip(w, state):
total += weight * feature
return total
def sample_action(probabilities, uniform_draw):
cumulative = 0.0
for action, probability in enumerate(probabilities):
cumulative += probability
if uniform_draw < cumulative:
return action
return len(probabilities) - 1
def train(env, n_features, n_actions, episodes=5_000, alpha_actor=0.01,
alpha_critic=0.05, gamma=0.99):
theta = [[0.0] * n_features for _ in range(n_actions)] # the actor
w = [0.0] * n_features # the critic
for _ in range(episodes):
state = env.reset()
done = False
# gamma^t on the step index. The derivation requires it and nearly
# every implementation silently drops it; writing it out here makes the
# omission elsewhere a decision rather than an accident.
discount_so_far = 1.0
while not done:
probabilities = policy_probabilities(theta, state)
action = sample_action(probabilities, random.random())
next_state, reward, done = env.step(action)
# Zero at a true terminal: there is no future to bootstrap from.
next_value = 0.0 if done else state_value(w, next_state)
value = state_value(w, state)
# One quantity, two jobs.
td_error = reward + gamma * next_value - value
# Critic. The gradient of (target - V)^2 with respect to w is
# -2 * td_error * dV/dw, and dV/dw is s for a linear critic. This
# is a semi-gradient: the target's own dependence on w, through
# V(s'), is not differentiated.
for index, feature in enumerate(state):
w[index] += alpha_critic * td_error * feature
# Actor. The score, weighted by the advantage.
#
# Note that td_error enters here as a NUMBER: the critic's
# parameters play no part in this update at all. That is the
# hand-written equivalent of detaching the advantage from the
# graph, and forgetting it in an autodiff implementation trains
# the critic to make advantages LARGE rather than accurate.
for candidate in range(n_actions):
indicator = 1.0 if candidate == action else 0.0
coefficient = discount_so_far * td_error * (indicator - probabilities[candidate])
for index, feature in enumerate(state):
theta[candidate][index] += alpha_actor * coefficient * feature
discount_so_far *= gamma
state = next_state
return theta, wO(|A| x d) per step for the actor plus O(d) for the critic, in the interpreter. One update per environment step, no batching.
How this language actually wants it written.
"""Actor-critic - shared trunk, GAE advantages, losses reported separately."""
from __future__ import annotations
from dataclasses import dataclass
from typing import NamedTuple
import torch
from torch import Tensor, nn
from torch.distributions import Categorical
class Rollout(NamedTuple):
states: Tensor # (T, N, features)
actions: Tensor # (T, N)
rewards: Tensor # (T, N)
values: Tensor # (T, N) critic output at collection time
# A true terminal has no future. A time-limit truncation does, and its
# bootstrap must be kept. One done flag conflates them, and the resulting
# policy is visibly short-sighted at the episode limit.
terminals: Tensor # (T, N)
truncations: Tensor # (T, N)
last_value: Tensor # (N,) bootstrap for the step after the rollout
@dataclass(frozen=True)
class ActorCriticConfig:
n_features: int
n_actions: int
trunk: tuple[int, ...] = (64, 64)
learning_rate: float = 3e-4
gamma: float = 0.99
# The bias-variance dial. Usually worth more attention than the learning
# rate: 0 is the fully bootstrapped advantage, 1 is Monte Carlo.
gae_lambda: float = 0.95
# The value regression has a much larger natural gradient scale than the
# policy surrogate, so on a shared trunk this is what stops the critic
# training the representation by itself.
value_coefficient: float = 0.5
entropy_coefficient: float = 0.01
grad_clip_norm: float = 0.5
def __post_init__(self) -> None:
if not 0.0 <= self.gamma < 1.0:
raise ValueError(f"gamma must be in [0, 1), got {self.gamma}")
if not 0.0 <= self.gae_lambda <= 1.0:
raise ValueError(f"lambda must be in [0, 1], got {self.gae_lambda}")
if self.value_coefficient <= 0.0:
raise ValueError("a non-positive value coefficient leaves the critic untrained")
class ActorCritic(nn.Module):
"""One trunk, two heads.
Features useful for valuing a state are useful for choosing in it, so the
representation is shared - which is also why the value coefficient exists.
"""
def __init__(self, config: ActorCriticConfig) -> None:
super().__init__()
layers: list[nn.Module] = []
width = config.n_features
for size in config.trunk:
layers += [nn.Linear(width, size), nn.Tanh()]
width = size
self.trunk = nn.Sequential(*layers)
self.policy_head = nn.Linear(width, config.n_actions)
self.value_head = nn.Linear(width, 1)
def forward(self, states: Tensor) -> tuple[Tensor, Tensor]:
features = self.trunk(states)
return self.policy_head(features), self.value_head(features).squeeze(-1)
def generalized_advantage(rollout: Rollout, config: ActorCriticConfig) -> tuple[Tensor, Tensor]:
"""GAE in one reversed pass: A_t = delta_t + gamma * lambda * A_(t+1).
The recursion is exact, so the exponentially weighted sum of TD errors
costs the same single backward pass a plain return would.
"""
horizon = rollout.rewards.shape[0]
advantages = torch.zeros_like(rollout.rewards)
running = torch.zeros_like(rollout.last_value)
next_value = rollout.last_value
for step in range(horizon - 1, -1, -1):
# Terminal: no future, so the bootstrap is zero and the accumulated
# advantage is cut. Truncated: the future exists, so the bootstrap is
# kept and only the accumulation is cut, because the trajectory the
# rollout saw does not continue.
keeps_future = (~rollout.terminals[step]).float()
continues = keeps_future * (~rollout.truncations[step]).float()
delta = (
rollout.rewards[step]
+ config.gamma * next_value * keeps_future
- rollout.values[step]
)
running = delta + config.gamma * config.gae_lambda * continues * running
advantages[step] = running
next_value = rollout.values[step]
# The critic regresses onto advantage + value, which is the same target the
# advantage was built from and keeps the two consistent.
returns = advantages + rollout.values
return advantages, returns
def explained_variance(predicted: Tensor, actual: Tensor) -> float:
"""1 - Var(actual - predicted) / Var(actual).
The single most informative number in an actor-critic run, and the one most
often not logged. Below zero the critic is worse than predicting the mean,
which means every advantage downstream of it is noise.
"""
variance = actual.var()
if variance == 0:
return 0.0
return float(1.0 - (actual - predicted).var() / variance)
def update(
model: ActorCritic,
optimizer: torch.optim.Optimizer,
rollout: Rollout,
config: ActorCriticConfig,
) -> dict[str, float]:
advantages, returns = generalized_advantage(rollout, config)
states = rollout.states.reshape(-1, config.n_features)
actions = rollout.actions.reshape(-1)
advantages = advantages.reshape(-1)
returns = returns.reshape(-1)
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
logits, values = model(states)
distribution = Categorical(logits=logits)
# detach() is load-bearing. The advantage is a WEIGHT, not a prediction; if
# it stays attached, the policy term backpropagates into the value head and
# trains the critic to make advantages large rather than correct. Nothing
# errors, both curves look plausible, and the objective is silently wrong.
policy_loss = -(distribution.log_prob(actions) * advantages.detach()).mean()
value_loss = nn.functional.mse_loss(values, returns)
entropy = distribution.entropy().mean()
loss = policy_loss + config.value_coefficient * value_loss - config.entropy_coefficient * entropy
optimizer.zero_grad(set_to_none=True)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip_norm)
optimizer.step()
# Reported separately, never summed. The policy surrogate's value means
# nothing and the value loss means something; adding them discards the only
# interpretable half and hides a degrading policy inside an improving critic.
return {
"policy_surrogate": float(policy_loss),
"value_loss": float(value_loss),
"entropy": float(entropy),
"explained_variance": explained_variance(values.detach(), returns),
}What changed: The one-step advantage becomes GAE, which is a continuous bias-variance dial computed in the same single reversed pass. Terminal and truncated become separate flags with different bootstrap handling, because conflating them teaches the agent the world ends at the step limit. The advantage is explicitly detached — left attached it trains the critic toward large advantages rather than accurate ones, with no error and plausible curves. The trunk is shared with a value coefficient balancing two gradient scales, and the losses are returned separately alongside explained variance, which is the number that actually diagnoses the run.
One shared forward per batch yielding both heads; GAE is one reversed pass over (T, N).
Faster, and the justification for each change.
"""A2C - environments in lockstep, GAE over the whole (T, N) block at once.
The name is the optimization. A3C ran asynchronous workers applying lock-free
updates from stale parameters; A2C steps N environments synchronously and
applies one update. The synchronous version is simpler and at least as good,
which localized where the benefit had actually come from: decorrelation across
independent environments, not the asynchrony.
"""
from __future__ import annotations
import numpy as np
import torch
from numpy.typing import NDArray
from torch import Tensor
class RolloutStore:
"""Pre-allocated (T, N) storage, sized once and written in place.
Appending per step and stacking at the boundary allocates T times and then
copies the whole block; a fixed buffer does neither, and the rollout is the
hot path of the whole method.
"""
def __init__(self, horizon: int, n_envs: int, n_features: int) -> None:
self.states = np.zeros((horizon, n_envs, n_features), dtype=np.float32)
self.actions = np.zeros((horizon, n_envs), dtype=np.int64)
self.rewards = np.zeros((horizon, n_envs), dtype=np.float32)
self.values = np.zeros((horizon, n_envs), dtype=np.float32)
self.terminals = np.zeros((horizon, n_envs), dtype=np.float32)
self.truncations = np.zeros((horizon, n_envs), dtype=np.float32)
# Reused every rollout rather than reallocated.
self.advantages = np.zeros((horizon, n_envs), dtype=np.float32)
def add(self, step: int, states, actions, rewards, values, terminals, truncations) -> None:
self.states[step] = states # in place, no allocation per step
self.actions[step] = actions
self.rewards[step] = rewards
self.values[step] = values
self.terminals[step] = terminals
self.truncations[step] = truncations
def gae_vectorized(
store: RolloutStore,
last_value: NDArray[np.float32],
gamma: float,
lam: float,
) -> NDArray[np.float32]:
"""GAE for every environment simultaneously.
The recursion is sequential in TIME and cannot be vectorized along it - but
it is fully independent across environments, so each step of the reversed
pass is one length-N array operation rather than N scalar ones. The Python
loop shrinks from T*N iterations to T.
"""
horizon = store.rewards.shape[0]
running = np.zeros_like(last_value)
next_value = last_value
for step in range(horizon - 1, -1, -1):
keeps_future = 1.0 - store.terminals[step]
continues = keeps_future * (1.0 - store.truncations[step])
# delta, the decay and the accumulation computed as one expression per
# step, writing straight into the pre-allocated advantage block.
delta = store.rewards[step] + gamma * next_value * keeps_future - store.values[step]
running = delta + gamma * lam * continues * running
store.advantages[step] = running
next_value = store.values[step]
return store.advantages
@torch.no_grad()
def act_and_value(model: torch.nn.Module, states: Tensor, generator: torch.Generator):
"""One shared forward yields BOTH heads.
Running the policy and the critic as separate passes duplicates the trunk -
which is most of the compute - for no reason. Sharing it means collection
costs one forward per step regardless of how many heads hang off it.
"""
logits, values = model(states)
# Gumbel-max: argmax(logits + gumbel noise) is already a softmax sample, so
# no normalization and no per-environment distribution object is built.
uniform = torch.rand(logits.shape, generator=generator, device=logits.device)
gumbel = -torch.log(-torch.log(uniform + 1e-20) + 1e-20)
actions = (logits + gumbel).argmax(dim=1)
return actions, values
def flatten_for_update(store: RolloutStore, advantages: NDArray[np.float32]):
"""Collapse (T, N) to (T*N,) with reshape and no copy.
The gradient is a sum over independent (state, action, advantage) triples,
so time-major or env-major makes no difference to the estimator and the
layout can be chosen purely for contiguity.
"""
n_features = store.states.shape[-1]
return (
torch.from_numpy(store.states.reshape(-1, n_features)),
torch.from_numpy(store.actions.reshape(-1)),
torch.from_numpy(advantages.reshape(-1)),
torch.from_numpy((advantages + store.values).reshape(-1)),
)What changed: A2C is itself the optimization: N environments step in lockstep and one synchronous update replaces A3C's asynchronous lock-free workers, which is simpler and at least as good — locating the benefit in decorrelation rather than asynchrony. The GAE recursion stays sequential in time, because it genuinely is, but becomes one length-N array operation per step instead of N scalar ones, shrinking the Python loop from T·N iterations to T. Storage is a (T, N) block sized once and written in place, and a single shared forward returns both heads rather than paying for the trunk twice.
Batch work to amortize interpreter overhead
N environments step in lockstep under one forward pass, so the per-step interpreter cost is paid once per batch rather than once per environment — this is what A2C is.
Trade-off: Environments must step synchronously, so a slow one stalls the batch, and reproducibility now depends on the environment count as well as the seed.
Eliminate Python-level loops over samples
Each step of the reversed GAE pass is one length-N array operation, cutting the Python loop from T·N iterations to T.
Trade-off: The time axis genuinely cannot be vectorized — the recursion depends on its own previous result — so T iterations remain and grow with rollout length.
Pre-allocate output arrays and use in-place operations
The (T, N) rollout block and the advantage buffer are sized once and written in place, so a rollout performs no allocation and no end-of-rollout stacking copy.
Trade-off: Horizon and environment count are fixed at construction, so changing either means rebuilding the store rather than adjusting a parameter.
Fuse adjacent operations so intermediates are never materialized
One shared trunk forward returns both heads, and Gumbel-max sampling skips normalization entirely, so neither the trunk nor the probability vector is computed twice.
Trade-off: Sampled probabilities are never formed, so anything needing them — a logged propensity, an importance ratio for a later PPO epoch — must recompute them.
One forward per batch step across all environments; GAE as T array operations instead of T x N scalar ones. Illustrative, not a measured benchmark.