Hands on RL 之 Off-policy Maximum Entropy Actor-Critic (SAC)

Hands on RL 之 Off-policy Maximum Entropy Actor-Critic (SAC)

​ Soft actor-critic方法又被称为Off-policy maximum entropy actor-critic algorithm。

1. 理论基础

1.1 Maximum Entropy Reinforcement Learning, MERL

​ 在MERL原论文中，引入了熵的概念，熵定义如下：

随机变量$x$符合$P$的概率分布，那么随机变量$x$的熵$\mathcal{H}(P)$为
$$\mathcal{H}(P)={\mathbb{E}}_{x\sim P}[-\log p(x)]$$
标准的RL算法的目标是能够找到最大化累计收益的策略
$${\pi }_{\text{std}}^{\ast }=\arg \underset{\pi }{\max }\sum _t{\mathbb{E}}_{(s_t,a_t)\sim {\rho }_{\pi }}[r(s_t,a_t)]$$
引入了熵最大化的RL算法的目标是
$${\pi }_{MERL}^{\ast }=\arg \underset{\pi }{\max }\sum _t{\mathbb{E}}_{(s_t,a_t)\sim {\rho }_{\pi }}[r(s_t,a_t)+\alpha \mathcal{H}(\pi (\cdot |s_t))]$$
其中，${\rho }_{\pi }$表示在策略$\pi$下状态空间对的概率分布，$\alpha$是温度系数，用于调节对熵的重视程度。

类似的，我们也可以在RL的动作值函数，和状态值函数中同样引入熵的概念。

标准的RL算法的值函数
$$\begin{array}{rl}\text{standard Q function:}{Q}^{\pi }(s,a)& ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t,a_t)|s_0=s,a_0=a]\\ \text{standard V function:}{V}^{\pi }(s)& ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t,a_t)|s_0=s]\end{array}$$

根据MERL的目标函数，在值函数中引入熵的概念之后就获得了Soft Value Function, SVF
$$\begin{array}{rl}\text{Soft Q function:}{Q}_{\text{soft}}^{\pi }(s,a)& ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t,a_t)+\alpha \sum _{t=1}^{\infty }{\gamma }^t\mathcal{H}(\pi (\cdot |s_t))|s_0=s,a_0=a]\\ \text{Soft V function:}{V}_{\text{soft}}^{\pi }(s)& ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[\sum _{t=0}^{\infty }{\gamma }^t(r(s_t,a_t)+\alpha \mathcal{H}(\pi (\cdot |s_t)))|s_0=s]\end{array}$$

观察上式我们就可以获得Soft Bellman Equation
$$\begin{gathered} \begin{array}{rl}{Q}_{\text{soft}}^{\pi }(s,a)& ={\mathbb{E}}_{s^{\prime }\sim p(s^{\prime }|s,a) \\ {a}^{\prime }\sim \pi }[r(s,a)+\gamma (Q_{\text{soft}}^{\pi }(s^{\prime },a^{\prime })+\alpha \mathcal{H}(\pi (\cdot |s^{\prime })))]\\ & ={\mathbb{E}}_{s^{\prime }\sim p(s^{\prime }|s,a)}[r(s,a)+\gamma V_{\text{soft}}^{\pi }(s)]\end{array} \end{gathered}$$

$$\begin{array}{rl}{V}_{\text{soft}}^{\pi }(s)& ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[\sum _{t=0}^{\infty }{\gamma }^t(r(s_t,a_t)+\alpha \mathcal{H}(\pi (\cdot |s_t)))|s_0=s]\\ & ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t,a_t)+\alpha \sum _{t=1}^{\infty }{\gamma }^t\mathcal{H}(\pi (\cdot |s_t))+\alpha \mathcal{H}(\pi (\cdot |s_t)|s_0=s]\\ & ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{\pi }}[Q_{\text{soft}}^{\pi }(s,a)|s_0=s,a_0=a]+\alpha \mathcal{H}(\pi (\cdot |s_t))\\ & ={\mathbb{E}}_{a\sim \pi }[Q_{\text{soft}}^{\pi }(s,a)-\alpha \log \pi (a|s)]\end{array}$$

1.2 Soft Policy Evaluation and Soft Policy Improvement in SAC

soft Q function的值迭代公式
$$\begin{gathered} \begin{array}{rl}{Q}_{\text{soft}}^{\pi }(s,a)& ={\mathbb{E}}_{s^{\prime }\sim p(s^{\prime }|s,a) \\ {a}^{\prime }\sim \pi }[r(s,a)+\gamma (Q_{\text{soft}}^{\pi }(s^{\prime },a^{\prime })+\alpha \mathcal{H}(\pi (\cdot |s^{\prime })))]\\ & ={\mathbb{E}}_{s^{\prime }\sim p(s^{\prime }|s,a)}[r(s,a)+\gamma V_{\text{soft}}^{\pi }(s)]\end{array}\text{\hspace{1em}(1.1)(1.2)} \end{gathered}$$

soft V function的值迭代公式

$$V_{\text{soft}}^{\pi }(s)={\mathbb{E}}_{a\sim \pi }[Q_{\text{soft}}^{\pi }(s,a)-\alpha \log \pi (a|s)]\text{\hspace{1em}(1.3)}$$

在SAC中如果我们只打算维持一个Q值函数，那么使用式子$(1,1)$进行值迭代即可，

如果需要同时维持Q，V两个函数，那么使用式子$(1,2),(1,3)$进行值迭代。

下面直接给出训练中的损失函数

V值函数的损失函数

$$J_{V(\psi )}={\mathbb{E}}_{s_t\sim \mathcal{D}}[\frac{1}{2}(V_{\psi }(s_t)-{\mathbb{E}}_{a_t\sim {\pi }_{\varphi }}[Q_{\theta }(s_t,a_t)-\log {\pi }_{\varphi }(a_t|s_t)]{)}^2]$$

Q值函数的损失函数

$$\begin{gathered} J_{Q(\theta )}={\mathbb{E}}_{(s_t,a_t)\sim \mathcal{D}}\left[\frac{1}{2}\right(Q_{\theta }(s_t,a_t)-\hat{Q}(s_t,a_t)\left)\right] \\ \hat{Q}(s_t,a_t)=r(s_t,a_t)+\gamma {\mathbb{E}}_{s_{t+1}\sim p}[V_{\psi }(s_{t+1})] \\ \hat{Q}(s_t,a_t)=r(s_t,a_t)+\gamma Q_{\theta }(s_{t+1},a_{t+1}) \end{gathered}$$

策略$\pi$的损失函数

$$J_{\pi }(\varphi )={\mathbb{E}}_{s_t\sim \mathcal{D}}\left[{\mathbb{D}}_{KL}\right({\pi }_{\varphi }(\cdot |s_t)|\left|\frac{\exp (Q_{\theta }(s_t,\cdot ))}{Z_{\theta }(s_t)}\right)]$$

其中，$Z_{\theta }(s_t)$是分配函数 (partition function)

KL散度的定义如下

​ 假设对随机变量$\xi$，存在两个概率分布P和Q，其中P是真实的概率分布，Q是较容易获得的概率分布。如果$\xi$是离散的随机变量，那么定义从P到Qd KL散度为

$${\mathbb{D}}_{KL}(P||Q)=\sum _{i}P(i)\ln (\frac{P(i)}{Q(i)})$$

如果$\xi$是连续变量，则定义从P到Q的KL散度为

$${\mathbb{D}}_{KL}(P||Q)={\int }_{-\infty }^{\infty }p(x)\ln (\frac{p(x)}{q(x)})dx$$

那么根据离散的KL散度定义，我们可以将策略$\pi$的损失函数展开如下

$$\begin{array}{rl}{J}_{\pi }(\varphi )& ={\mathbb{E}}_{s_t\sim \mathcal{D}}\left[{\mathbb{D}}_{KL}\right({\pi }_{\varphi }(\cdot |s_t)|\left|\frac{\exp (Q_{\theta }(s_t,\cdot ))}{Z_{\theta }(s_t)}\right)]\\ & ={\mathbb{E}}_{s_t\sim \mathcal{D}}[\sum {\pi }_{\varphi }(\cdot |s_t)\ln (\frac{{\pi }_{\varphi }(\cdot |s_t)Z_{\theta }(s_t)}{\exp (Q_{\theta }(s_t,\cdot ))})]\\ & ={\mathbb{E}}_{s_t\sim \mathcal{D},a_t\sim {\pi }_{\varphi }}[\ln ({\pi }_{\varphi }(\cdot |s_t)+\ln (Z_{\theta }(s_t))-Q_{\theta }(s_t,\cdot )]\\ & ={\mathbb{E}}_{s_t\sim \mathcal{D},a_t\sim {\pi }_{\varphi }}[\ln ({\pi }_{\varphi }(\cdot |s_t))-Q_{\theta }(s_t,\cdot )]\end{array}\text{\hspace{1em}(1.4)}$$

最后一步是因为，分配函数$Z_{\theta }$与策略$\pi$无关，所以对梯度没有影响，所以在目标函数中可以直接忽略。

然后又引入了 reparameter的方法，

$$a_t=f_{\varphi }({ϵ}_t;s_t),ϵ\sim \mathcal{N}$$

将上式带入$(1.4)$中则有

$$J_{\pi }(\varphi )={\mathbb{E}}_{s_t\sim \mathcal{D},a_t\sim {\pi }_{\varphi }}\left[\ln \right({\pi }_{\varphi }(f_{\varphi }({ϵ}_t;s_t)|s_t\left)\right)-Q_{\theta }(s_t,f_{\varphi }({ϵ}_t;s_t))]$$

最后，只需要不断收集数据，缩小这两个损失函数，就可以得到收敛到一个解。

1.3 Two Q Value Neural Network

​ SAC的策略网络与DDPG中策略网络直接估计动作值大小不同，SAC的策略网络是估计连续动作空间Gaussian分布的均值和方差，再从这个Gaussian 分布中均匀采样获得连续的动作值。SAC中可以维护一个V值函数网络，也可以不维护V值函数网络，但是SAC中必须的是两个Q值函数网络及两个Q值网络的target network和一个策略网络$\pi$。为什么要维护两个Q值网络呢？这是因为为了减少Q值网络存在的过高估计的问题。然后用两个Q值网络中较小的Q值网络来计算损失函数，假设存在两个Q值网络$Q({\theta }_1),Q({\theta }_2)$和其对应的target network $Q({\theta }_1^{-}),Q({\theta }_2^{-})$

​ 那么Q值函数的损失函数就成了

$$\begin{array}{rl}{J}_Q(\theta )& ={\mathbb{E}}_{(s_t,a_t)\sim \mathcal{D}}\left[\frac{1}{2}\right(Q_{\theta }(s_t,a_t)-\hat{Q}(s_t,a_t)\left)\right]\\ \hat{Q}(s_t,a_t)& =r(s_t,a_t)+\gamma V_{\psi }(s_{t+1})\\ V_{\psi }(s_{t+1})& =Q_{\theta }(s_{t+1},a_{t+1})-\alpha \log (\pi (a_{t+1}|s_{t+1}))\end{array}\text{\hspace{1em}(1.5)(1.6)(1.7)}$$

结合式子$1.5,1.6,1.7$则有

$$J_Q(\theta )={\mathbb{E}}_{(s_t,a_t)\sim \mathcal{D}}\left[\frac{1}{2}\right(Q_{\theta }(s_t,a_t)-(r(s_t,a_t)+\gamma (\underset{j=1,2}{\min }{Q}_{{\theta }_j^{-}}(s_{t+1},a_{t+1})-\alpha \log \pi (a_{t+1}|s_{t+1}))\left){)}^2\right]$$

因为SAC也是一种离线的算法，为了让训练更加稳定，这里使用了目标Q网络$Q({\theta }^{-})$，同样是两个目标网络，与两个Q网络相对应，SAC中目标Q网络的更新方式和DDPG中一致。

​ 之前推导了策略$\pi (\varphi )$的损失函数如下

$$J_{\pi }(\varphi )={\mathbb{E}}_{s_t\sim \mathcal{D},a_t\sim {\pi }_{\varphi }}[\ln ({\pi }_{\varphi }(\cdot |s_t))-Q_{\theta }(s_t,\cdot )]$$

再引入reparameter化简之后，就成了

$$J_{\pi }(\varphi )={\mathbb{E}}_{s_t\sim \mathcal{D},a_t\sim {\pi }_{\varphi }}\left[\ln \right({\pi }_{\varphi }(f_{\varphi }({ϵ}_t;s_t)|s_t\left)\right)-Q_{\theta }(s_t,f_{\varphi }({ϵ}_t;s_t))]$$

利用双Q网络取最小，再将上述目标函数进行改写

$$J_{\pi }(\varphi )={\mathbb{E}}_{s_t\sim \mathcal{D},a_t\sim {\pi }_{\varphi }}\left[\ln \right({\pi }_{\varphi }(f_{\varphi }({ϵ}_t;s_t)|s_t\left)\right)-\underset{j=1,2}{\min }{Q}_{{\theta }_j}(s_t,f_{\varphi }({ϵ}_t;s_t))]$$

1.4 Tricks

​ SAC还集成了一些别的算法的常用技巧，比如Replay Buffer来产生independent and identically distribution的样本，使用了Double DQN中target network的思想使用两个网络。SAC中最重要的一个技巧是自动调整熵正则项，在 SAC 算法中，如何选择熵正则项的系数非常重要。在不同的状态下需要不同大小的熵：在最优动作不确定的某个状态下，熵的取值应该大一点；而在某个最优动作比较确定的状态下，熵的取值可以小一点。为了自动调整熵正则项，SAC 将强化学习的目标改写为一个带约束的优化问题：

$$\underset{\pi }{\max }{\mathbb{E}}_{\pi }[\sum _{t}r(s_t,a_t)]\text{s.t.}{\mathbb{E}}_{(s_t,a_t)\sim {\rho }_{\pi }}[-\log ({\pi }_t(a_t|s_t))]\ge {\mathcal{H}}_0$$

${\mathcal{H}}_0$ 是预先定义好的最小策略熵的阈值。通过一些数学技巧简化，可以得到温度$\alpha$的损失函数

$$J(\alpha )={\mathbb{E}}_{s_t\sim R,a_t\sim \pi (\cdot |s_t)}[-\alpha \log ({\pi }_t(a_t|s_t))-\alpha {\mathcal{H}}_0]$$

1.5 Pesudocode

1.6 What does code actually do

# Actor
class PolicyNet(nn.Module):
    """
    func->
        动作空间符合高斯正态分布,输出动作空间的均值mu,和标准差std
    """
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):
        super(PolicyNet, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=state_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(in_features=hidden_dim, out_features=action_dim)
        self.fc_std = nn.Sequential(
            nn.Linear(in_features=hidden_dim, out_features=action_dim),
            nn.Softplus()
        )
        self.action_bound = action_bound

    def forward(self, s):
        x = self.fc1(s)
        mu = self.fc_mu(x)
        std = self.fc_std(x)
        distribution = Normal(mu, std)
        normal_sample = distribution.rsample()
        # calculate the log probability density function(pdf)
        normal_log_prob = distribution.log_prob(normal_sample)
        # get action limit to [-1,1]
        action = torch.tanh(normal_sample)
        # get tanh_normal log pdf
        tanh_log_prob = normal_log_prob - torch.log(1 - action.pow(2) + 1e-7)
        # get action bounded
        action = action * self.action_bound
        return action, tanh_log_prob

​ 我们使用Policy Network来获得输出动作的mean、std, 然后通过Normal Distribution来获得该mean、std下的正态高斯分布，按照这个概率分布来进行一次采样即可以获得在这个概率分布下的采样值，然后我们利用$\tanh$函数将这个采样值限制在$(-1,1)$之间。然后计算该采样值的一个对数概率密度函数，我们已知正态高斯分布的概率密度函数为
$$f(x)=\frac{1}{\sqrt{2\pi }\delta }\exp (-\frac{(x-\mu {)}^2}{2{\delta }^2})$$
然后对这个概率密度函数取对数，这即是log_prob在计算的东西
$$\log (f(x))=-\frac{(x-\mu {)}^2}{2{\delta }^2}-\log \delta -\log (\sqrt{2\pi })$$
我们已知随机变量$X$概率密度函数为正态分布的概率密度函数$f(x)$（如上），求随机变量$Y=\tanh (X)$的概率密度函数，
$$f(y)=f(x)\times \frac{dx}{dy}$$
已知$y=\tanh (x)$，我们可以方便地求出$\frac{dy}{dx}$，再利用反函数求导法则就可以获得$\frac{dx}{dy}$
$$\frac{dy}{dx}={\tanh }^{\prime }(x)=1-{\tanh }^2(x)$$
利用反函数求导法则
$$\begin{array}{rl}x& =\phi (y)\\ y& =\psi (x)\\ {\phi }^{\prime }(y)& =\frac{1}{{\psi }^{\prime }(x)}=\frac{1}{\psi (\phi (y))}\end{array}$$
那么则有
$$\frac{dx}{dy}=\frac{1}{1-{\tanh }^2(x)}=\frac{1}{1-{\tanh }^2({\tanh }^{-1}(y))}=\frac{1}{1-y^2}$$
那么随机变量$Y$的概率密度函数为
$$\begin{array}{rl}f(y)& =f(x)\times \frac{dx}{dy}\\ & =f(x)\frac{1}{1-y^2}\\ \log f(y)& =\log f(x)-\log (1-y^2)\end{array}$$
推导完毕。

2. 代码实现

2.1 SAC处理连续动作空间

采用gymnasium中Pendulum-v1的连续动作环境，整体代码如下

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal

from tqdm import tqdm
import collections
import random
import numpy as np
import matplotlib.pyplot as plt
import gym

# replay buffer
class ReplayBuffer():
    def __init__(self, capacity):
        self.buffer = collections.deque(maxlen=capacity)
    
    def add(self, s, r, a, s_, d):
        self.buffer.append((s,r,a,s_,d))
    
    def sample(self, batch_size):
        transitions = random.sample(self.buffer, batch_size)
        states, rewards, actions, next_states, dones = zip(*transitions)
        return np.array(states), rewards, actions, np.array(next_states), dones
    
    def size(self):
        return len(self.buffer)

# Actor
class PolicyNet_Continuous(nn.Module):
    """动作空间符合高斯分布,输出动作空间的均值mu,和标准差std"""
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound):
        super(PolicyNet_Continuous, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=state_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(in_features=hidden_dim, out_features=action_dim)
        self.fc_std = nn.Sequential(
            nn.Linear(in_features=hidden_dim, out_features=action_dim),
            nn.Softplus()
        )
        self.action_bound = action_bound

    def forward(self, s):
        x = self.fc1(s)
        mu = self.fc_mu(x)
        std = self.fc_std(x)
        distribution = Normal(mu, std)
        normal_sample = distribution.rsample()
        normal_log_prob = distribution.log_prob(normal_sample)
        # get action limit to [-1,1]
        action = torch.tanh(normal_sample)
        # get tanh_normal log probability
        tanh_log_prob = normal_log_prob - torch.log(1 - action.pow(2) + 1e-7)
        # get action bounded
        action = action * self.action_bound
        return action, tanh_log_prob


# Critic
class QValueNet_Continuous(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(QValueNet_Continuous, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=state_dim + action_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(in_features=hidden_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.fc_out = nn.Linear(in_features=hidden_dim, out_features=1)
    
    def forward(self, s, a):
        cat = torch.cat([s,a], dim=1)
        x = self.fc1(cat)
        x = self.fc2(x)
        return self.fc_out(x)

# maximize entropy deep reinforcement learning SAC
class SAC_Continuous():
    def __init__(self, state_dim, hidden_dim, action_dim, action_bound,
                    actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma,
                    device):
        # actor
        self.actor = PolicyNet_Continuous(state_dim, hidden_dim, action_dim, action_bound).to(device)
        # two critics
        self.critic1 = QValueNet_Continuous(state_dim, hidden_dim, action_dim).to(device)
        self.critic2 = QValueNet_Continuous(state_dim, hidden_dim, action_dim).to(device)
        # two target critics
        self.target_critic1 = QValueNet_Continuous(state_dim, hidden_dim, action_dim).to(device)
        self.target_critic2 = QValueNet_Continuous(state_dim, hidden_dim, action_dim).to(device)
        # initialize with same parameters
        self.target_critic1.load_state_dict(self.critic1.state_dict())
        self.target_critic2.load_state_dict(self.critic2.state_dict())
        # specify optimizers
        self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)
        self.optimizer_critic1 = torch.optim.Adam(self.critic1.parameters(), lr=critic_lr)
        self.optimizer_critic2 = torch.optim.Adam(self.critic2.parameters(), lr=critic_lr)
        # 使用alpha的log值可以使训练稳定
        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float, requires_grad = True)
        self.optimizer_log_alpha = torch.optim.Adam([self.log_alpha], lr=alpha_lr)

        self.target_entropy = target_entropy
        self.gamma = gamma
        self.tau = tau
        self.device = device
    
    def take_action(self, state):
        state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
        action, _ = self.actor(state)
        return [action.item()]
    
    # calculate td target
    def calc_target(self, rewards, next_states, dones):
        next_action, log_prob = self.actor(next_states)
        entropy = -log_prob
        q1_values = self.target_critic1(next_states, next_action)
        q2_values = self.target_critic2(next_states, next_action)
        next_values = torch.min(q1_values, q2_values) + self.log_alpha.exp() * entropy
        td_target = rewards + self.gamma * next_values * (1-dones)
        return td_target

    # soft update method
    def soft_update(self, net, target_net):
        for param_target, param in zip(target_net.parameters(), net.parameters()):
            param_target.data.copy_(param_target.data * (1.0-self.tau) + param.data * self.tau)
        
    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1,1).to(self.device)
        actions = torch.tensor(transition_dict['actions'], dtype=torch.float).view(-1,1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1,1).to(self.device)

        rewards = (rewards + 8.0) / 8.0     #对倒立摆环境的奖励进行重塑，方便训练

        # update two Q-value network
        td_target = self.calc_target(rewards, next_states, dones).detach()
        critic1_loss = torch.mean(F.mse_loss(td_target, self.critic1(states, actions)))
        critic2_loss = torch.mean(F.mse_loss(td_target, self.critic2(states, actions)))

        self.optimizer_critic1.zero_grad()
        critic1_loss.backward()
        self.optimizer_critic1.step()
        self.optimizer_critic2.zero_grad()
        critic2_loss.backward()
        self.optimizer_critic2.step()

        # update policy network
        new_actions, log_prob = self.actor(states)
        entropy = - log_prob
        q1_value = self.critic1(states, new_actions)
        q2_value = self.critic2(states, new_actions)
        actor_loss = torch.mean(-self.log_alpha.exp() * entropy - torch.min(q1_value, q2_value))
        self.optimizer_actor.zero_grad()
        actor_loss.backward()
        self.optimizer_actor.step()

        # update temperature alpha
        alpha_loss = torch.mean((entropy - self.target_entropy).detach() * self.log_alpha.exp())
        self.optimizer_log_alpha.zero_grad()
        alpha_loss.backward()
        self.optimizer_log_alpha.step()

        # soft update target Q-value network
        self.soft_update(self.critic1, self.target_critic1)
        self.soft_update(self.critic2, self.target_critic2)


def train_off_policy_agent(env, agent, num_episodes, replay_buffer, minimal_size, batch_size, render, seed_number):
    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes/10), desc='Iteration %d'%(i+1)) as pbar:
            for i_episode in range(int(num_episodes/10)):
                observation, _ = env.reset(seed=seed_number)
                done = False
                episode_return = 0

                while not done:
                    if render:
                        env.render()
                    action = agent.take_action(observation)
                    observation_, reward, terminated, truncated, _ = env.step(action)
                    done = terminated or truncated
                    replay_buffer.add(observation, action, reward, observation_, done)
                    # swap states
                    observation = observation_
                    episode_return += reward
                    if replay_buffer.size() > minimal_size:
                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)
                        transition_dict = {
                            'states': b_s,
                            'actions': b_a,
                            'rewards': b_r,
                            'next_states': b_ns,
                            'dones': b_d
                        }
                        agent.update(transition_dict)
                return_list.append(episode_return)
                if(i_episode+1) % 10 == 0:
                    pbar.set_postfix({
                        'episode': '%d'%(num_episodes/10 * i + i_episode + 1),
                        'return': "%.3f"%(np.mean(return_list[-10:]))
                    })
                pbar.update(1)
    env.close()
    return return_list

def moving_average(a, window_size):
    cumulative_sum = np.cumsum(np.insert(a, 0, 0)) 
    middle = (cumulative_sum[window_size:] - cumulative_sum[:-window_size]) / window_size
    r = np.arange(1, window_size-1, 2)
    begin = np.cumsum(a[:window_size-1])[::2] / r
    end = (np.cumsum(a[:-window_size:-1])[::2] / r)[::-1]
    return np.concatenate((begin, middle, end))

def plot_curve(return_list, mv_return, algorithm_name, env_name):
    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list, c='gray', alpha=0.6)
    plt.plot(episodes_list, mv_return)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('{} on {}'.format(algorithm_name, env_name))
    plt.show()



if __name__ == "__main__":

    # reproducible
    seed_number = 0
    random.seed(seed_number)
    np.random.seed(seed_number)
    torch.manual_seed(seed_number)

    num_episodes = 150     # episodes length
    hidden_dim = 128        # hidden layers dimension
    gamma = 0.98            # discounted rate
    actor_lr = 1e-4         # lr of actor
    critic_lr = 1e-3        # lr of critic
    alpha_lr = 1e-4
    tau = 0.005             # soft update parameter
    buffer_size = 10000
    minimal_size = 1000
    batch_size = 64

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    env_name = 'Pendulum-v1'

    render = False
    if render:
        env = gym.make(id=env_name, render_mode='human')
    else:
        env = gym.make(id=env_name)

    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]  
    action_bound = env.action_space.high[0]
    # entropy初始化为动作空间维度的负数
    target_entropy = - env.action_space.shape[0]

    replaybuffer = ReplayBuffer(buffer_size)
    agent = SAC_Continuous(state_dim, hidden_dim, action_dim, action_bound, actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma, device)
    return_list = train_off_policy_agent(env, agent, num_episodes, replaybuffer, minimal_size, batch_size, render, seed_number)

    mv_return = moving_average(return_list, 9)
    plot_curve(return_list, mv_return, 'SAC', env_name)

所获得的回报曲线如下图

2.2 SAC处理离散动作空间

采用gaymnasium中的CartPole-v1离散动作空间的环境。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical

from tqdm import tqdm
import collections
import random
import numpy as np
import matplotlib.pyplot as plt
import gym

# replay buffer
class ReplayBuffer():
    def __init__(self, capacity):
        self.buffer = collections.deque(maxlen=capacity)
    
    def add(self, s, r, a, s_, d):
        self.buffer.append((s,r,a,s_,d))
    
    def sample(self, batch_size):
        transitions = random.sample(self.buffer, batch_size)
        states, rewards, actions, next_states, dones = zip(*transitions)
        return np.array(states), rewards, actions, np.array(next_states), dones
    
    def size(self):
        return len(self.buffer)

# Actor
class PolicyNet_Discrete(nn.Module):
    """动作空间服从离散的概率分布,输出每个动作的概率值"""
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet_Discrete, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=state_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.fc2 = nn.Sequential(
            nn.Linear(in_features=hidden_dim, out_features=action_dim),
            nn.Softmax(dim=1)
        )

    def forward(self, s):
        x = self.fc1(s)
        return self.fc2(x)


# Critic
class QValueNet_Discrete(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(QValueNet_Discrete, self).__init__()
        self.fc1 = nn.Sequential(
            nn.Linear(in_features=state_dim, out_features=hidden_dim),
            nn.ReLU()
        )
        self.fc2 = nn.Linear(in_features=hidden_dim, out_features=action_dim)
    
    def forward(self, s):
        x = self.fc1(s)
        return self.fc2(x)

# maximize entropy deep reinforcement learning SAC
class SAC_Continuous():
    def __init__(self, state_dim, hidden_dim, action_dim,
                    actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma,
                    device):
        # actor
        self.actor = PolicyNet_Discrete(state_dim, hidden_dim, action_dim).to(device)
        # two critics
        self.critic1 = QValueNet_Discrete(state_dim, hidden_dim, action_dim).to(device)
        self.critic2 = QValueNet_Discrete(state_dim, hidden_dim, action_dim).to(device)
        # two target critics
        self.target_critic1 = QValueNet_Discrete(state_dim, hidden_dim, action_dim).to(device)
        self.target_critic2 = QValueNet_Discrete(state_dim, hidden_dim, action_dim).to(device)
        # initialize with same parameters
        self.target_critic1.load_state_dict(self.critic1.state_dict())
        self.target_critic2.load_state_dict(self.critic2.state_dict())
        # specify optimizers
        self.optimizer_actor = torch.optim.Adam(self.actor.parameters(), lr=actor_lr)
        self.optimizer_critic1 = torch.optim.Adam(self.critic1.parameters(), lr=critic_lr)
        self.optimizer_critic2 = torch.optim.Adam(self.critic2.parameters(), lr=critic_lr)
        # 使用alpha的log值可以使训练稳定
        self.log_alpha = torch.tensor(np.log(0.01), dtype=torch.float, requires_grad = True)
        self.optimizer_log_alpha = torch.optim.Adam([self.log_alpha], lr=alpha_lr)

        self.target_entropy = target_entropy
        self.gamma = gamma
        self.tau = tau
        self.device = device
    
    def take_action(self, state):
        state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
        probs = self.actor(state)
        action_dist = Categorical(probs)
        action = action_dist.sample()
        return action.item()
    
    # calculate td target
    def calc_target(self, rewards, next_states, dones):
        next_probs = self.actor(next_states)
        next_log_probs = torch.log(next_probs + 1e-8)
        entropy = -torch.sum(next_probs * next_log_probs, dim=1, keepdim=True)
        q1_values = self.target_critic1(next_states)
        q2_values = self.target_critic2(next_states)
        min_qvalue = torch.sum(next_probs * torch.min(q1_values, q2_values),
                               dim=1,
                               keepdim=True)
        next_value = min_qvalue + self.log_alpha.exp() * entropy
        td_target = rewards + self.gamma * next_value * (1 - dones)
        return td_target

    # soft update method
    def soft_update(self, net, target_net):
        for param_target, param in zip(target_net.parameters(), net.parameters()):
            param_target.data.copy_(param_target.data * (1.0-self.tau) + param.data * self.tau)
        
    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'], dtype=torch.float).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1,1).to(self.device)
        actions = torch.tensor(transition_dict['actions'], dtype=torch.int64).view(-1,1).to(self.device)
        next_states = torch.tensor(transition_dict['next_states'], dtype=torch.float).to(self.device)
        dones = torch.tensor(transition_dict['dones'], dtype=torch.float).view(-1,1).to(self.device)

        rewards = (rewards + 8.0) / 8.0     #对倒立摆环境的奖励进行重塑，方便训练

        # update two Q-value network
        td_target = self.calc_target(rewards, next_states, dones).detach()
        critic1_loss = torch.mean(F.mse_loss(td_target, self.critic1(states).gather(dim=1,index=actions)))
        critic2_loss = torch.mean(F.mse_loss(td_target, self.critic2(states).gather(dim=1,index=actions)))

        self.optimizer_critic1.zero_grad()
        critic1_loss.backward()
        self.optimizer_critic1.step()
        self.optimizer_critic2.zero_grad()
        critic2_loss.backward()
        self.optimizer_critic2.step()

        # update policy network
        probs = self.actor(states)
        log_probs = torch.log(probs + 1e-8)
        entropy = -torch.sum(probs * log_probs, dim=1, keepdim=True)
        q1_value = self.critic1(states)
        q2_value = self.critic2(states)
        min_qvalue = torch.sum(probs * torch.min(q1_value, q2_value),
                               dim=1,
                               keepdim=True)  # 直接根据概率计算期望
        actor_loss = torch.mean(-self.log_alpha.exp() * entropy - min_qvalue)
        self.optimizer_actor.zero_grad()
        actor_loss.backward()
        self.optimizer_actor.step()

        # update temperature alpha
        alpha_loss = torch.mean((entropy - self.target_entropy).detach() * self.log_alpha.exp())
        self.optimizer_log_alpha.zero_grad()
        alpha_loss.backward()
        self.optimizer_log_alpha.step()

        # soft update target Q-value network
        self.soft_update(self.critic1, self.target_critic1)
        self.soft_update(self.critic2, self.target_critic2)


def train_off_policy_agent(env, agent, num_episodes, replay_buffer, minimal_size, batch_size, render, seed_number):
    return_list = []
    for i in range(10):
        with tqdm(total=int(num_episodes/10), desc='Iteration %d'%(i+1)) as pbar:
            for i_episode in range(int(num_episodes/10)):
                observation, _ = env.reset(seed=seed_number)
                done = False
                episode_return = 0

                while not done:
                    if render:
                        env.render()
                    action = agent.take_action(observation)
                    observation_, reward, terminated, truncated, _ = env.step(action)
                    done = terminated or truncated
                    replay_buffer.add(observation, action, reward, observation_, done)
                    # swap states
                    observation = observation_
                    episode_return += reward
                    if replay_buffer.size() > minimal_size:
                        b_s, b_a, b_r, b_ns, b_d = replay_buffer.sample(batch_size)
                        transition_dict = {
                            'states': b_s,
                            'actions': b_a,
                            'rewards': b_r,
                            'next_states': b_ns,
                            'dones': b_d
                        }
                        agent.update(transition_dict)
                return_list.append(episode_return)
                if(i_episode+1) % 10 == 0:
                    pbar.set_postfix({
                        'episode': '%d'%(num_episodes/10 * i + i_episode + 1),
                        'return': "%.3f"%(np.mean(return_list[-10:]))
                    })
                pbar.update(1)
    env.close()
    return return_list

def moving_average(a, window_size):
    cumulative_sum = np.cumsum(np.insert(a, 0, 0)) 
    middle = (cumulative_sum[window_size:] - cumulative_sum[:-window_size]) / window_size
    r = np.arange(1, window_size-1, 2)
    begin = np.cumsum(a[:window_size-1])[::2] / r
    end = (np.cumsum(a[:-window_size:-1])[::2] / r)[::-1]
    return np.concatenate((begin, middle, end))

def plot_curve(return_list, mv_return, algorithm_name, env_name):
    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list, c='gray', alpha=0.6)
    plt.plot(episodes_list, mv_return)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.title('{} on {}'.format(algorithm_name, env_name))
    plt.show()



if __name__ == "__main__":

    # reproducible
    seed_number = 0
    random.seed(seed_number)
    np.random.seed(seed_number)
    torch.manual_seed(seed_number)

    num_episodes = 200    # episodes length
    hidden_dim = 128        # hidden layers dimension
    gamma = 0.98            # discounted rate
    actor_lr = 1e-3         # lr of actor
    critic_lr = 1e-2        # lr of critic
    alpha_lr = 1e-2
    tau = 0.005             # soft update parameter
    buffer_size = 10000
    minimal_size = 500
    batch_size = 64


    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    env_name = 'CartPole-v1'

    render = False
    if render:
        env = gym.make(id=env_name, render_mode='human')
    else:
        env = gym.make(id=env_name)

    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n  

    # entropy初始化为-1
    target_entropy = -1

    replaybuffer = ReplayBuffer(buffer_size)
    agent = SAC_Continuous(state_dim, hidden_dim, action_dim, actor_lr, critic_lr, alpha_lr, target_entropy, tau, gamma, device)
    return_list = train_off_policy_agent(env, agent, num_episodes, replaybuffer, minimal_size, batch_size, render, seed_number)

    mv_return = moving_average(return_list, 9)
    plot_curve(return_list, mv_return, 'SAC', env_name)

Reference

Hands on RL

SAC(Soft Actor-Critic)阅读笔记

Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor