用PyTorch实现SAC：最大熵Actor-Critic

用PyTorch实现SAC：最大熵Actor-Critic

Soft Actor-Critic (SAC) 是近几年在连续控制任务里用得比较多的离线策略算法。它把策略熵直接放进了目标函数，让智能体在最大化奖励的同时保持一定的随机性，探索能力比很多传统方法强。这篇文章写一下它的核心原理，再给出一个完整的 PyTorch 实现。

从标准RL到最大熵目标

传统强化学习的目标是最大化累计折扣奖励：

$$J(\pi) = \mathbb{E}{\pi} \left[ \sum{t=0}^T \gamma^t r(s_t, a_t) \right]$$

SAC 在做一件不一样的事：它在奖励后面多加了一个熵项，变成

$$J(\pi) = \mathbb{E}{\pi} \left[ \sum{t=0}^T \gamma^t \left( r(s_t, a_t) + \alpha \mathcal{H}(\pi(\cdot|s_t)) \right) \right]$$

其中 $\mathcal{H}(\pi(\cdot|s_t)) = -\mathbb{E}_{a \sim \pi} [\log \pi(a|s_t)]$，就是策略在那时的熵。$\alpha$ 控制熵的权重。

这么改的好处很直接：如果你只最大化奖励，策略可能会过早变得确定，再也不去碰那些看起来不怎么样的动作，但可能错过更好的长期路径。加了这个熵项之后，策略会被推着保持一定随机性，探索更多，训练也更稳定。

网络结构和更新方式

SAC 用了 Actor-Critic 那一套，但 Critic 部分有两个 Q 网络（$Q_{\theta_1}, Q_{\theta_2}$），训练时取两者中的较小值作为目标，这主要是为了压住 Q 值过估计。另外它还用目标网络做软更新，让价值估计平滑一些。

Q 网络的目标值按 Bellman 方程算：

$$y = r + \gamma (1 - \text{done}) \cdot V_{\psi'}(s')$$

损失就是 MSE：

$$J_Q = \mathbb{E}{(s, a, r, s') \sim D} \left[ \left( Q{\theta_i}(s, a) - y \right)^2 \right]$$

策略网络这边输出的是均值和标准差，通过重参数化采样得到动作，再用 $\tanh$ 把动作压到 $[-1, 1]$ 再乘上 max_action。对数概率的计算要带上 $\tanh$ 的修正项，这是 SAC 论文里的标准操作。策略更新的目标是最小化：

$$J_\pi = \mathbb{E}{s \sim D, a \sim \pi} \left[ \alpha \log \pi\phi(a|s) - \min_{i=1,2} Q_{\theta_i}(s, a) \right]$$

直觉上就是让动作的 Q 值尽可能高，同时不要变得太确定（熵尽量大）。

代码实现

下面是一个干净的 PyTorch 实现，环境用的是 Pendulum-v1。代码分成了参数配置、网络定义、回放缓冲区和训练循环几个部分，可以直接跑。

参数配置

"""SAC 实现 - 2024.12"""
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import gym
import random
from collections import deque

GAMMA = 0.99
TAU = 0.005           # 目标网络软更新系数
ALPHA = 0.2           # 熵系数，视环境可能需要调
LR = 0.001
BATCH_SIZE = 256
MEMORY_CAPACITY = 100000

策略网络

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, 256)
        self.fc2 = nn.Linear(256, 256)
        self.mean = nn.Linear(256, action_dim)
        self.log_std = nn.Linear(256, action_dim)
        self.max_action = max_action

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        mean = self.mean(x)
        log_std = self.log_std(x).clamp(-20, 2)
        std = torch.exp(log_std)
        return mean, std

    def sample(self, state):
        mean, std = self.forward(state)
        normal = torch.distributions.Normal(mean, std)
        x_t = normal.rsample()
        y_t = torch.tanh(x_t)
        action = y_t * self.max_action

        # 带 tanh 修正的对数概率
        log_prob = normal.log_prob(x_t)
        log_prob -= torch.log(1 - y_t.pow(2) + 1e-6)
        log_prob = log_prob.sum(dim=-1, keepdim=True)
        return action, log_prob

Q 网络

class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(QNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim + action_dim, 256)
        self.fc2 = nn.Linear(256, 256)
        self.fc3 = nn.Linear(256, 1)

    def forward(self, state, action):
        x = torch.cat([state, action], dim=-1)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

经验回放缓冲区

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)

    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return (np.array(states), np.array(actions), np.array(rewards),
                np.array(next_states), np.array(dones))

    def __len__(self):
        return len(self.buffer)

SAC 智能体

class SACAgent:
    def __init__(self, state_dim, action_dim, max_action):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.actor = PolicyNetwork(state_dim, action_dim, max_action).to(self.device)
        self.q1 = QNetwork(state_dim, action_dim).to(self.device)
        self.q2 = QNetwork(state_dim, action_dim).to(self.device)

        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=LR)
        self.q1_optimizer = optim.Adam(self.q1.parameters(), lr=LR)
        self.q2_optimizer = optim.Adam(self.q2.parameters(), lr=LR)

        self.replay_buffer = ReplayBuffer(MEMORY_CAPACITY)
        self.max_action = max_action

    def select_action(self, state):
        state = torch.FloatTensor(state).to(self.device).unsqueeze(0)
        action, _ = self.actor.sample(state)
        return action.cpu().detach().numpy()[0]

    def train(self):
        if len(self.replay_buffer) < BATCH_SIZE:
            return

        states, actions, rewards, next_states, dones = self.replay_buffer.sample(BATCH_SIZE)
        states = torch.FloatTensor(states).to(self.device)
        actions = torch.FloatTensor(actions).to(self.device)
        rewards = torch.FloatTensor(rewards).unsqueeze(1).to(self.device)
        next_states = torch.FloatTensor(next_states).to(self.device)
        dones = torch.FloatTensor(dones).unsqueeze(1).to(self.device)

        # 更新两个 Q 网络
        with torch.no_grad():
            next_actions, log_probs = self.actor.sample(next_states)
            target_q1 = self.q1(next_states, next_actions)
            target_q2 = self.q2(next_states, next_actions)
            target_q = torch.min(target_q1, target_q2) - ALPHA * log_probs
            q_target = rewards + GAMMA * (1 - dones) * target_q

        q1_loss = ((self.q1(states, actions) - q_target) ** 2).mean()
        q2_loss = ((self.q2(states, actions) - q_target) ** 2).mean()

        self.q1_optimizer.zero_grad()
        q1_loss.backward()
        self.q1_optimizer.step()

        self.q2_optimizer.zero_grad()
        q2_loss.backward()
        self.q2_optimizer.step()

        # 更新策略网络
        new_actions, log_probs = self.actor.sample(states)
        q1_new = self.q1(states, new_actions)
        q2_new = self.q2(states, new_actions)
        q_new = torch.min(q1_new, q2_new)
        actor_loss = (ALPHA * log_probs - q_new).mean()

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

    def update_replay_buffer(self, state, action, reward, next_state, done):
        self.replay_buffer.push(state, action, reward, next_state, done)

主训练循环

env = gym.make("Pendulum-v1")
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.shape[0]
max_action = float(env.action_space.high[0])

agent = SACAgent(state_dim, action_dim, max_action)
num_episodes = 500

for episode in range(num_episodes):
    state = env.reset()
    episode_reward = 0
    done = False
    while not done:
        action = agent.select_action(state)
        next_state, reward, done, _ = env.step(action)
        agent.update_replay_buffer(state, action, reward, next_state, done)
        agent.train()
        state = next_state
        episode_reward += reward
    print(f"Episode {episode}, Reward: {episode_reward}")

调参这些事

上面的超参数在 Pendulum 上能跑出一个不错的结果，但换到其他环境可能需要微调。最敏感的通常是熵系数 $\alpha$：如果环境奖励尺度很大，$\alpha$ 可能也需要相应放大，否则熵项的影响几乎可以忽略。反过来，如果环境本身随机性很强，$\alpha$ 设得太大会让策略太散，收敛慢。

SAC 的后续版本里其实有自动调节 $\alpha$ 的机制，让它在训练过程中自适应，不过这里为了简洁只给了固定版本。另外，学习率和缓冲区大小也值得一试：样本复杂的环境往往需要更大的回放缓冲区，学习率太高可能让 Q 网络震荡。

双 Q 网络确实能有效抑制过估计，但也不是万能药；如果训练还是不稳，可以试试稍微增大 TAU 让目标网络更新慢一点。

运行环境建议：Python 3.11.5 + torch 2.1.0 + gym 0.26.2。

参考文献

• Haarnoja, Tuomas, et al. "Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor." arXiv preprint arXiv:1801.01290 (2018).
• Haarnoja, Tuomas, et al. "Soft actor-critic algorithms and applications." arXiv preprint arXiv:1812.05905 (2018).