1. Maximum Entropy Reinforcement Learning Objective

Standard reinforcement learning maximizes expected return. Maximum Entropy Reinforcement Learning augments immediate rewards with the policy’s Shannon entropy $\mathcal{H}(\pi(\cdot \mid s_t))$ to encourage robust exploration and multimodal action discovery:

$$\mathcal{J}(\pi) = \sum_{t=0}^\infty \gamma^t \mathbb{E}_{(s_t, a_t) \sim \rho_\pi}\left[ r(s_t, a_t) + \alpha \mathcal{H}\big(\pi(\cdot \mid s_t)\big) \right]$$

where $\alpha > 0$ is the temperature parameter governing exploration strength.

Soft State-Value & Action-Value Definitions

Under entropy regularization, the Soft State-Value Function incorporates expected log-densities:

$$V(s) = \mathbb{E}_{a \sim \pi}\left[ Q(s, a) - \alpha \log \pi(a \mid s) \right]$$

Defining the Soft Bellman Evaluation Operator $\mathcal{T}^\pi$:

$$\left( \mathcal{T}^\pi Q \right)(s, a) = r(s, a) + \gamma \mathbb{E}_{s' \sim \mathcal{P}}\left[ V(s') \right]$$

2. Soft Policy Evaluation Contraction Proof

Theorem (Soft Bellman Contraction)

Consider the Soft Bellman evaluation operator $\mathcal{T}^\pi: \mathbb{R}^{|\mathcal{S}| \times |\mathcal{A}|} \to \mathbb{R}^{|\mathcal{S}| \times |\mathcal{A}|}$. For any stationary policy $\pi$, $\mathcal{T}^\pi$ is a $\gamma$-contraction in the supremum norm.

Proof

Let $Q_1$ and $Q_2$ be two action-value estimators. Using the soft Bellman operator:

$$\left\| \mathcal{T}^\pi Q_1 - \mathcal{T}^\pi Q_2 \right\|_\infty = \max_{s, a} \left| \gamma \mathbb{E}_{s' \sim \mathcal{P}}\left[ \mathbb{E}_{a' \sim \pi}\left[ Q_1(s', a') - \alpha \log \pi(a' \mid s') - \left( Q_2(s', a') - \alpha \log \pi(a' \mid s') \right) \right] \right] \right|$$

Notice that the policy entropy terms $\alpha \log \pi(a' \mid s')$ cancel out exactly:

$$\left\| \mathcal{T}^\pi Q_1 - \mathcal{T}^\pi Q_2 \right\|_\infty = \gamma \max_{s, a} \left| \mathbb{E}_{s' \sim \mathcal{P}, \, a' \sim \pi}\left[ Q_1(s', a') - Q_2(s', a') \right] \right|$$

Applying the supremum norm bound across expected states and actions:

$$\left\| \mathcal{T}^\pi Q_1 - \mathcal{T}^\pi Q_2 \right\|_\infty \le \gamma \max_{s', a'} \left| Q_1(s', a') - Q_2(s', a') \right| = \gamma \|Q_1 - Q_2\|_\infty \quad \blacksquare$$

By Banach’s Fixed-Point Theorem, soft policy evaluation converges uniquely to exact soft action-values $Q^\pi$.


3. Reparameterization Trick & Continuous Squashed Policies

In continuous action spaces $\mathcal{A} \subset \mathbb{R}^d$, Soft Actor-Critic (SAC) (Haarnoja et al., 2018) parameterizes Gaussian policies $\pi_\phi(\mathbf{a} \mid \mathbf{s})$ bounded to $[-1, 1]^d$ via an invertible hyperbolic tangent squash function $\mathbf{a} = \tanh(\mathbf{u})$.

Reparameterization & Change-of-Variables Log-Density Jacobian

To differentiate through stochastic sampling, pre-squash actions $\mathbf{u}_t$ are reparameterized using isotropic Gaussian noise $\boldsymbol{\epsilon}_t \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$:

$$\mathbf{u}_t = \boldsymbol{\mu}_\phi(\mathbf{s}_t) + \boldsymbol{\sigma}_\phi(\mathbf{s}_t) \odot \boldsymbol{\epsilon}_t, \quad \mathbf{a}_t = \tanh(\mathbf{u}_t)$$

By the probability change-of-variables formula, the exact log-density of squashed action $\mathbf{a}_t$ requires subtracting the log-determinant of the Jacobian $\frac{\partial \mathbf{a}}{\partial \mathbf{u}} = \text{diag}(1 - \tanh^2(\mathbf{u}))$:

$$\log \pi_\phi(\mathbf{a}_t \mid \mathbf{s}_t) = \log \mathcal{N}\left(\mathbf{u}_t; \; \boldsymbol{\mu}_\phi(\mathbf{s}_t), \boldsymbol{\sigma}_\phi^2(\mathbf{s}_t)\right) - \sum_{i=1}^d \log\left( 1 - \tanh^2(u_{t, i}) + \epsilon \right)$$

Enabling low-variance path-wise analytical policy gradient updates:

$$\nabla_\phi J_\pi(\phi) = \nabla_\phi \mathbb{E}_{\boldsymbol{\epsilon}_t \sim \mathcal{N}}\left[ \alpha \log \pi_\phi\left( f_\phi(\boldsymbol{\epsilon}_t; \mathbf{s}_t) \;\big|\; \mathbf{s}_t \right) - Q_\theta\left( \mathbf{s}_t, f_\phi(\boldsymbol{\epsilon}_t; \mathbf{s}_t) \right) \right]$$

4. Automatic Dual Temperature Optimization

Fixed temperature $\alpha$ requires manual hyperparameter tuning per environment. SAC automates entropy control by reformulating maximum entropy RL as constrained optimization against a target minimum entropy threshold $\bar{\mathcal{H}} = -d$:

$$\mathcal{L}(\alpha) = \mathbb{E}_{\mathbf{s}_t \sim \mathcal{D}, \, \mathbf{a}_t \sim \pi_\phi}\left[ -\alpha \left( \log \pi_\phi(\mathbf{a}_t \mid \mathbf{s}_t) + \bar{\mathcal{H}} \right) \right]$$

When entropy drops below $\bar{\mathcal{H}}$, gradient updates automatically increase $\alpha$ to force exploratory behavior.


5. Python Verification: Squashed Gaussian Action & Log-Density Jacobian

import numpy as np

class SquashedGaussianPolicy:
    """Rigorous numpy calculation of SAC squashed Gaussian action sampling and exact log-prob Jacobian."""
    def __init__(self, action_dim: int = 3):
        self.dim = action_dim

    def sample(self, mu: np.ndarray, log_std: np.ndarray) -> Tuple[np.ndarray, float]:
        """mu shape: (dim,), log_std shape: (dim,)"""
        std = np.exp(log_std)
        eps = np.random.normal(0, 1, size=self.dim)
        
        # Reparameterized pre-squash Gaussian action u
        u = mu + std * eps
        
        # Squashed action in [-1, 1]^d
        action = np.tanh(u)
        
        # Base normal log probability density
        logp_normal = -0.5 * np.sum(((u - mu) / (std + 1e-8))**2 + 2 * log_std + np.log(2 * np.pi))
        
        # Change-of-variables Jacobian correction factor
        log_det_jacobian = np.sum(np.log(1.0 - action**2 + 1e-6))
        
        exact_log_prob = float(logp_normal - log_det_jacobian)
        return action, exact_log_prob

if __name__ == "__main__":
    np.random.seed(42)
    policy = SquashedGaussianPolicy(action_dim=2)
    sample_mu = np.array([0.5, -1.2])
    sample_log_std = np.array([-0.3, -0.7])
    
    a, logp = policy.sample(sample_mu, sample_log_std)
    print("Sampled Squashed Continuous Action (Bounded in [-1, 1]^2):")
    print(" ", np.round(a, 4))
    print(f"Exact Log-Probability Density with Jacobian Correction: {logp:.4f}")