class SquashedGaussianPolicy(nnx.Module):
"""A state-dependent Gaussian squashed through a = c tanh(u)."""
def __init__(self, obs_dim, act_dim, hidden=64, rngs=None):
self.trunk = nnx.Sequential(
nnx.Linear(obs_dim, hidden, rngs=rngs), jax.nn.relu,
nnx.Linear(hidden, hidden, rngs=rngs), jax.nn.relu)
self.mu = nnx.Linear(hidden, act_dim, rngs=rngs)
self.log_std = nnx.Linear(hidden, act_dim, rngs=rngs)
def __call__(self, obs):
h = self.trunk(obs)
return self.mu(h), jnp.exp(jnp.clip(self.log_std(h), -5, 2))
def log_prob(self, u, mean, std):
"""log pi at a = c tanh(u), from the pre-squash u the sampler keeps."""
logdet = 2 * (jnp.log(2.0) - u - jax.nn.softplus(-2 * u))
return (jax.scipy.stats.norm.logpdf(u, mean, std)
- logdet - jnp.log(c)).sum(-1)
def sample(self, obs, key):
"""A reparameterized action and its log-probability, differentiable."""
mean, std = self(obs)
u = mean + std * jax.random.normal(key, std.shape)
return c * jnp.tanh(u), self.log_prob(u, mean, std)
def act(self, obs, rng):
if not hasattr(self, '_fwd'): # compile the fixed-shape acting
self._fwd = nnx.cached_partial(nnx.jit(lambda net, o: net(o)),
self) # forward, once
mean, std = self._fwd(jnp.asarray(obs))
u = np.asarray(mean) + np.asarray(std) * rng.standard_normal(
mean.shape, dtype=np.float32)
return c * np.tanh(u)
def act_greedy(self, obs, rng=None):
mean, _ = self(jnp.asarray(obs))
return c * np.tanh(np.asarray(mean))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
@nnx.jit
def sac_step(agent, obs, act, rew, next_obs, term, key):
"""One SAC update: soft critic regression, pathwise actor step, Polyak.
Fixed batch shapes, so the step compiles once (:numref:`sec_compilation`)."""
k1, k2 = jax.random.split(key)
a2, logp2 = agent.actor.sample(next_obs, k1) # fresh, live policy
y = rew + gamma * (1 - term) * (
agent.min_q(next_obs, a2, agent.targets) - alpha * logp2)
def q_loss(qs):
x = jnp.concatenate([obs, act], -1)
return sum(((q(x).squeeze(-1) - y) ** 2).mean() for q in qs)
_, grads = nnx.value_and_grad(q_loss)(agent.qs)
agent.opt_q.update(agent.qs, grads)
def pi_loss(actor):
a, logp = actor.sample(obs, k2)
return (alpha * logp - agent.min_q(obs, a)).mean(), logp
(_, logp), grads = nnx.value_and_grad(pi_loss, has_aux=True)(agent.actor)
agent.opt_pi.update(agent.actor, grads)
for q, tnet in zip(agent.qs, agent.targets): # Polyak: the drifting copy
nnx.update(tnet, jax.tree.map(lambda p, tp: tau * p + (1 - tau) * tp,
nnx.state(q, nnx.Param),
nnx.state(tnet, nnx.Param)))
return logp.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: [7400, 5400, 6400]
single critic: env steps to a trailing five-episode average of -200: [6400, 5400, 7200]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.18 0.07 -0.12]
single critic: entropy over the last 20 episodes, per seed [ 0.01 0.15 -0.03]- 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 -151.5, delivered (soft) -112.1, gap -39.4, plain -118.1
SAC, seed 1: promised -115.0, delivered (soft) -83.5, gap -31.5, plain -93.3
SAC, seed 2: promised -184.2, delivered (soft) -136.7, gap -47.5, plain -141.3
single critic, seed 0: promised -112.1, delivered (soft) -109.8, gap -2.3, plain -118.8
single critic, seed 1: promised -80.5, delivered (soft) -82.8, gap +2.4, plain -93.1
single critic, seed 2: promised -162.6, delivered (soft) -134.2, gap -28.4, plain -141.8Neither 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.