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

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 - torch.tanh(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