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
maximum-entropy control · pathwise actor gradients · twin critics · squashed Gaussian actions
The Maximum-Entropy Objective
:numref:sec_regularized’s KL penalty, with a uniform reference, applied at every 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]
The entropy term is part of the objective, so it changes the optimum and can make the final policy stochastic. Continuous actions: differential entropy, can be negative. SAC = the off-policy actor-critic of this objective :cite:Haarnoja.Zhou.Abbeel.ea.2018.
Soft Policy Evaluation
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.2962661899Soft Policy Improvement
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 Squashed Gaussian Policy
Section 14.7 allowed the environment to clip torque because the score estimator did not differentiate through actions. SAC’s pathwise estimator does:
- outside the interval the clip’s derivative is zero, so samples beyond a boundary provide no actor gradient
- 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]
Numerical Stability near the Boundary
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.6137The guarded expression remains at -13.8155 after saturation, so it no longer represents the density accurately. Quadrature checks normalization without training:
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 Components of SAC
- twin critics, minimum: actor optimization can favor positive critic errors (15.4’s maximization argument); the minimum of two independent critics is pessimistic :cite:
Fujimoto.vanHoof.Meger.2018 - Polyak targets: w^- \leftarrow \tau w + (1-\tau) w^-, with half-life \ln 2 / \tau \approx 139 updates; the target changes continuously
- no target actor: \tilde{a}' from the live policy; the stochastic policy smooths its own targets
- no ratios: the target is independent of the collecting policy, and the actor samples new actions; the buffer shifts the state distribution (:numref:
sec_offline) ReplayBufferC: one column widened to float vectors
The SAC 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 impose an incorrect zero continuation value at step 200.
Sample Efficiency
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 actor gradients and replay both improve reuse of collected transitions. The two variants are indistinguishable on this axis.
Policy Entropy
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]- decreases during rapid return improvement and finishes near zero; differential entropy can be negative, and the policy remains stochastic
- 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 differ by about ten points on this task
Critic Calibration
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 variant substantially overestimates the realized soft return. The single critic is near calibrated; the minimum of two critics underestimates by thirty to sixty points on every seed. At this budget, the return curves remain similar.