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 }_{}}[\sum _{t=0}^{\infty }{c}^{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 }_{}}[\sum _{t=0}^{\infty }{c}^{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 }_{}}[\sum _{t=0}^{\infty }{c}^{t}r(s_t,a_t)+a\sum _{t=1}^{\infty }{c}^t\mathcal{H}(\pi (\cdot |s_t))|s_0=s,a_0=a]\\ \text{Soft V function:}{V}_{\text{soft}}^{}(s)& =E_{s_t,a_t\sim {\rho }_{}}\left[\sum _{t=0}^{\infty }{c}^t\right(r(s_t,a_t)+\alpha H\left(\pi \right(\cdot |s_t)))|s_0=s]\end{array}$$

Determine the specific orientation of the Soft Bellman Equation
$$\begin{gathered} \begin{array}{rl}{Q}_{\text{soft}}^{}(s,a)& ={\mathbb{E}}_{s^{\prime }\sim p(s^{\prime }|s,a) \\ {a}^{\prime }\sim \pi }[r(s,a)+c(Q_{\text{soft}}^{}(s^{\prime },a^{\prime })+\alpha H\left(\pi \right(\cdot |s^{\prime })))]\\ & ={\mathbb{E}}_{s^{\prime }\sim p(s^{\prime }|s,a)}\left[r\right(s,a)+\gamma V_{\text{soft}}^{}(s)]\end{array} \end{gathered}$$

$$\begin{array}{rl}{V}_{\text{soft}}^{}(s)& ={\mathbb{E}}_{s_t,a_t\sim {\rho }_{}}[\sum _{t=0}^{\infty }{c}^t(r(s_t,a_t)+\alpha H\left(\pi \right(\cdot |s_t)))|s_0=s]\\ & =E_{s_t,a_t\sim {\rho }_{}}[\sum _{t=0}^{\infty }{c}^{t}r(s_t,a_t)+a\sum _{t=1}^{\infty }{c}^t\mathcal{H}(\pi (\cdot |s_t))+\alpha H\left(\pi \right(\cdot |s_t)|s_0=s]\\ & =E_{s_t,a_t\sim {\rho }_{}}[Q_{\text{soft}}^{}(s,a)|s_0=s,a_0=a]+\alpha H(\pi (\cdot |s_t))\\ & =E_{a\sim \pi }[Q_{\text{soft}}^{}(s,a)-a\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}}^{}(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_{}(s_t,a_t)-\hat{Q}(s_t,a_t)\left)\right] \\ \hat{Q}(s_t,a_t)=r(s_t,a_t)+c{E}_{s_{t+1}\sim p}[V_{}(s_{t+1})] \\ \hat{Q}(s_t,a_t)=r(s_t,a_t)+\gamma Q_{}(s_{t+1},a_{t+1}) \end{gathered}$$

Strategy π \piLoss function of$\pi$

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

Among them, $Z_{}\left(s_t\right)$ is the partition function

The definition of KL divergence is as follows

​ Suppose that for random variables $\xi$ , there are two probability distributions P and Q, where P is the real probability distribution and Q is the easier to obtain probability distribution. If$\xi$ is a discrete random variable, then define the KL divergence from P to Qd as

$${\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}_{}\left(\varphi \right)& ={\mathbb{E}}_{s_t\sim \mathcal{D}}\left[{\mathbb{D}}_{KL}\right(p_{}(\cdot |s_t)|\left|\frac{\exp (Q_{}(s_t,\cdot ))}{Z_{}(s_t)}\right)]\\ & =E_{s_t\sim D}[\sum {Pi}_{}(\cdot |s_t)\ln (\frac{{Pi}_{}(\cdot |s_t)Z_{}(s_t)}{\exp \left(Q_{}\right(s_t,\cdot ))})]\\ & =E_{s_t\sim \mathcal{D},a_t\sim {\pi }_{}}\left[\ln \right(p_{}(\cdot |s_t)+\ln (Z_{}(s_t))-Q_{}(s_t,\cdot )]\\ & =E_{s_t\sim D,a_t\sim {\pi }_{}}\left[ln\right(p_{}(\cdot |s_t))-Q_{}(s_t,\cdot )]\end{array}\text{\hspace{1em}(1.4)}$$

The last step is because, the distribution function $Z_{}$With strategy $\pi$ has nothing to do, so it has no impact on the gradient, so it can be directly ignored in the objective function.

Then the reparameter method was introduced,

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

Put the above equation into $In\; (1.4)$ there is

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

Finally, you only need to continue to collect data and shrink these two loss functions to converge to a solution.

1.3 Two Q Value Neural Network

The policy network of SAC is different from the policy network in DDPG that directly estimates the size of the action value. The policy network of SAC estimates the mean and variance of the Gaussian distribution in the continuous action space, and then uniformly samples the continuous action value from this Gaussian distribution. SAC can maintain a V-value function network or not maintain a V-value function network, but what is required in SAC are two Q-value function networks and the target network of the two Q-value networks and a policy network π \ $\pi$ . Why maintain two Q-value networks? This is to reduce the problem of overestimation in the Q-value network. Then use the smaller Q-value network of the two Q-value networks to calculate the loss function, assuming there are two Q-value networks$Q({\theta }_1),Q({\theta }_2)$ and its corresponding target network$Q\left({\theta }_1^{-}\right),Q({\theta }_2^{-})$

​Then the loss function of the Q-value function becomes

$$\begin{array}{rl}{J}_Q\left(i\right)& ={\mathbb{E}}_{(s_t,a_t)\sim \mathcal{D}}\left[\frac{1}{2}\right(Q_{}(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_{}\left(s_{t+1}\right)\\ V_{}(s_{t+1})& =Q_{}(s_{t+1},a_{t+1})-a\log \left(\pi \right(a\_{}_{t+1}|s_{t+1}))\end{array}\text{\hspace{1em}(1.5)(1.6)(1.7)}$$

Combining the formulas $1.5,1.6,1.7$ has

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

Because SAC is also an offline algorithm, in order to make the training more stable, the target Q network $Q({\theta }^{-})$, are also two target networks, corresponding to two Q networks. The update method of the target Q network in SAC is consistent with that in DDPG.

​ The strategy π ( ϕ ) \pi(\phi) was derived beforeThe loss function of$\pi$$($$\varphi$$)\; is\; as\; follows$

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

After introducing reparameter for simplification, it becomes

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

Use the double Q network to find the minimum, and then rewrite the above objective function

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

1.4 Tricks

SAC also integrates some common techniques of other algorithms, such as Replay Buffer to generate independent and identically distributed samples, and uses the idea of ​​target network in Double DQN to use two networks. One of the most important skills in SAC is to automatically adjust the entropy regularization term. In the SAC algorithm, how to choose the coefficient of the entropy regularization term is very important. Different states require different amounts of entropy: in a state where the optimal action is uncertain, the value of entropy should be larger; and in a state where the optimal action is relatively certain, the value of entropy can be small. a little. In order to automatically adjust the entropy regularization term, SAC rewrites the goal of reinforcement learning into a constrained optimization problem:

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

$H_0$is the predefined minimum policy entropy threshold. By simplifying with some mathematical techniques, we can get the temperature α \alphaLoss function of$\alpha$

$$J\left(a\right)={\mathbb{E}}_{s_t\sim R,a_t\sim \pi (\cdot |s_t)}[-alog(p_t(a_t|s_t))-a{H}_0]$$

1.5 Pesudocode

2. Code implementation

2.1 SAC handles continuous action space

Using the gymnasiumcontinuous Pendulum-v1action environment, the overall code is as follows

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)

The return curve obtained is as follows:

2.2 SAC deals with discrete action space

An environment employing gaymnasiumdiscrete CartPole-v1action spaces.

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) reading notes

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