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.2962661899

Soft 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]

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())

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.6137

The 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 without

The 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

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 impose an incorrect zero continuation value at step 200.

Sample Efficiency

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 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.03  0.04  0.07]
single critic: entropy over the last 20 episodes, per seed [-0.06  0.12  0.24]
  • 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  -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.3

Neither 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.