class SquashedGaussianPolicy(nn.Module):
"""A state-dependent Gaussian squashed through a = c tanh(u)."""
def __init__(self, obs_dim, act_dim, hidden=64):
super().__init__()
self.trunk = nn.Sequential(nn.Linear(obs_dim, hidden), nn.ReLU(),
nn.Linear(hidden, hidden), nn.ReLU())
self.mu = nn.Linear(hidden, act_dim)
self.log_std = nn.Linear(hidden, act_dim)
def forward(self, obs):
h = self.trunk(obs)
return self.mu(h), self.log_std(h).clamp(-5, 2).exp()
def log_prob(self, u, mean, std):
"""log pi at a = c tanh(u), from the pre-squash u the sampler keeps."""
logdet = 2 * (np.log(2) - u - nn.functional.softplus(-2 * u))
return (torch.distributions.Normal(mean, std).log_prob(u)
- logdet - np.log(c)).sum(-1)
def sample(self, obs):
"""A reparameterized action and its log-probability, differentiable."""
mean, std = self(obs)
u = mean + std * torch.randn_like(std)
return c * torch.tanh(u), self.log_prob(u, mean, std)
def act(self, obs, rng):
with torch.no_grad():
mean, std = self(torch.as_tensor(obs))
u = mean.numpy() + std.numpy() * rng.standard_normal(
mean.shape, dtype=np.float32)
return c * np.tanh(u)
def act_greedy(self, obs, rng=None):
with torch.no_grad():
return c * np.tanh(self(torch.as_tensor(obs))[0].numpy())Soft Actor-Critic
Dive into Deep Learning · §15.5
Soft Actor-Critic
the objective from 15.3 · the gradient from 14.7 · the critics from 15.4 · one new line of calculus
The Objective, Already Proved
:numref:sec_regularized’s KL penalty, uniform reference, charged per step; \alpha for \beta, the field’s convention:
J(\pi) = E_{\pi}\Big[ \sum_t \gamma^t \big( r_t + \alpha\, H(\pi(\cdot \mid s_t)) \big) \Big]
Not a bonus bolted onto training: part of the objective, so the optimum itself is stochastic, by design. Continuous actions: differential entropy, can be negative. SAC = the off-policy actor-critic of this objective :cite:Haarnoja.Zhou.Abbeel.ea.2018.
Soft Evaluation: One New Term
y = r + \gamma\, \big(1 - \mathbf{1}(s' \textrm{ terminal})\big) \Big( \min_{j=1,2} Q_{w_j^-}(s', \tilde{a}') - \alpha \log \pi_\theta(\tilde{a}' \mid s') \Big)
- s, a, r, s' from the buffer; \tilde{a}' fresh from the live policy: the expectation is under \pi_\theta
- PPO’s entropy bonus lived in the actor loss; this one lives in the critic target: the critic values entropy collected later
At the tilted optimum the bracket is 15.3’s logsumexp, exactly:
alpha = 0.1: E[Q - alpha log pi*] = 0.6415156557, alpha logsumexp(Q/alpha) = 0.6415156557
alpha = 0.5: E[Q - alpha log pi*] = 0.9890764345, alpha logsumexp(Q/alpha) = 0.9890764345
alpha = 2.0: E[Q - alpha log pi*] = 3.2962661899, alpha logsumexp(Q/alpha) = 3.2962661899Improvement Is the Proposition
15.3’s proof line, read with r \to Q(s, \cdot), uniform reference:
E_{\pi}[Q] + \alpha H(\pi) \;=\; \alpha \log Z - \alpha\, D_{\textrm{KL}}\big(\pi \,\Vert\, e^{Q/\alpha}/Z\big)
Maximizing the left and projecting onto the family in KL are the same optimization; they differ by \alpha \log Z, which does not depend on \theta.
The gradient through \tilde{a}_\theta(s, z) = c \tanh(\mu_\theta + \sigma_\theta z) is 14.7’s pathwise estimator, on a critic differentiable in a by construction.
Proposition (soft policy improvement): exact per-state maximization raises V everywhere; five lines from 15.3.
A Policy That Fits in a Box
14.7 let the environment clip the torque. The score estimator never differentiated through the action; the pathwise one does:
- outside the box the clip’s derivative is zero: no signal at the boundary, where swing-up lives
- a clipped Gaussian is not a density: atoms at \pm 2, and \alpha \log \pi is undefined on an atom
\log \pi(a \mid s) = \sum_i \Big[ \log \mathcal{N}(u_i; \mu_i, \sigma_i) - \log\big(1 - \tanh^2 u_i\big) - \log c \Big]
The Epsilon That Hides
1 - \tanh^2 u = 4 e^{-2u}/(1 + e^{-2u})^2, so \log(1 - \tanh^2 u) = 2(\log 2 - u - \operatorname{softplus}(-2u)), exact. What the guard + 1e-6 does instead:
u naive guarded stable
0 0.0000 0.0000 0.0000
3 -4.6187 -4.6186 -4.6187
8 -14.5561 -13.4256 -14.6137
10 -inf -13.8155 -18.6137
20 -inf -13.8155 -38.6137-13.8155 forever: the density stops charging for saturation and no curve looks wrong. And the whole change of variables is checkable by quadrature, at zero training cost:
mu = 0.0, sigma = 0.5: integrates to 1.000000 with the log-det, 1.653 without
mu = 0.7, sigma = 0.8: integrates to 1.000000 with the log-det, 1.132 withoutThe Machinery, Off the Shelf
- twin critics, min: the actor climbs the critic, a maximizer over its errors (15.4’s argument, argmax \to gradient ascent); the min of two independent critics is the cheapest pessimism :cite:
Fujimoto.vanHoof.Meger.2018 - Polyak targets: w^- \leftarrow \tau w + (1-\tau) w^-, half-life \ln 2 / \tau \approx 139 updates; drifts, never jumps
- no target actor: \tilde{a}' from the live policy; the stochastic policy smooths its own targets
- no ratios: the target never mentions the collector, the actor re-samples; the buffer shifts only the state distribution (:numref:
sec_offline) ReplayBufferC: one column widened to float vectors
One Update
def sac_step(agent, batch):
"""One SAC update: soft critic regression, pathwise actor step, Polyak."""
obs, act = torch.as_tensor(batch.obs), torch.as_tensor(batch.act)
rew, term = torch.as_tensor(batch.rew), torch.as_tensor(batch.term)
next_obs = torch.as_tensor(batch.next_obs)
with torch.no_grad(): # the critic target is data
a2, logp2 = agent.actor.sample(next_obs)
y = rew + gamma * (1 - term) * (
agent.min_q(next_obs, a2, agent.targets) - alpha * logp2)
x = torch.cat([obs, act], -1)
loss_q = sum(((q(x).squeeze(-1) - y) ** 2).mean() for q in agent.qs)
agent.opt_q.zero_grad()
loss_q.backward()
agent.opt_q.step()
a, logp = agent.actor.sample(obs) # fresh, from the live policy
loss_pi = (alpha * logp - agent.min_q(obs, a)).mean()
agent.opt_pi.zero_grad()
loss_pi.backward()
agent.opt_pi.step()
with torch.no_grad(): # Polyak: the drifting copy
for q, tnet in zip(agent.qs, agent.targets):
for pq, pt in zip(q.parameters(), tnet.parameters()):
pt.mul_(1 - tau).add_(tau * pq)
return float(logp.detach().mean())Pendulum has no terminal state: term is identically zero, the bootstrap is always taken; storing done would teach the agent that the world ends at step 200.
Under Ten Thousand Steps
SAC: env steps to a trailing five-episode average of -200: [6200, 7600, 7000]
single critic: env steps to a trailing five-episode average of -200: [6000, 6800, 5600]14.7’s REINFORCE spent 480{,}000 steps on this task and never reached -200. Pathwise gradient \times replay: more signal per sample, hundreds of updates per sample. The two arms are indistinguishable on this axis.
The Entropy the Policy Keeps
SAC: entropy over the last 20 episodes, per seed [-0.03 0.04 0.07]
single critic: entropy over the last 20 episodes, per seed [-0.06 0.12 0.24]- spent fast during the climb, overshooting below zero, partly bought back once the task is mastered; ends near zero: stochastic at convergence
- within about a nat of autotuning’s target \bar{H} = -\dim \mathcal{A} = -1 :cite:
Haarnoja.Zhou.Hartikainen.ea.2018; \alpha is an exchange rate in reward per nat, not a learning rate - deterministic and stochastic evaluations agree within about ten points: the noise was kept where it is cheap
Honest Promises
SAC, seed 0: promised -175.1, delivered (soft) -117.9, gap -57.1, plain -126.4
SAC, seed 1: promised -127.0, delivered (soft) -89.7, gap -37.3, plain -98.1
SAC, seed 2: promised -177.1, delivered (soft) -135.6, gap -41.6, plain -141.9
single critic, seed 0: promised -119.6, delivered (soft) -113.0, gap -6.6, plain -119.3
single critic, seed 1: promised -86.5, delivered (soft) -81.8, gap -4.7, plain -93.7
single critic, seed 2: promised -144.2, delivered (soft) -129.9, gap -14.3, plain -138.3Neither arm meaningfully over-promises. The single critic is near calibrated; the min under-promises by thirty to sixty points, on every seed, at identical policy quality: pessimism, measured, for free. On harder tasks that margin is what stands between this loop and 15.4’s self-confirming collapse.