策略梯度是一种基于策略的算法，相比于DQN一类的基于价值的算法，它会直接显式的学习一个目标策略。梯度下降的基础知识可以参考之前的博客强化学习（六）策略梯度和《动手学强化学习》部分内容。

Reference

[1] 《动手学强化学习》 https://hrl.boyuai.com/
[2] David Silver: https://www.youtube.com/watch?v=KHZVXao4qXs&t=4609s

我们假设目标策略${\pi }_{\theta }(a|s)$是一种随机策略，并且处处可微，$\theta$为对应参数。可以通过神经网络或线性模型对目标策略进行建模。输入某个状态，输出动作的概率分布。
（1）目标函数
我们期望获得一个最优策略，能够最大化策略在环境中的期望回报。
$$J(\theta )=E[V^{{\pi }_{\theta }}(s)]$$

根据贝尔曼方程，我们可以用Q函数表示目标函数
$$J(\theta )=\sum _{s\in S}d(s)\sum _{a\in A}{\pi }_{\theta }(a|s)Q^{{\pi }_{\theta }}(s,a)$$

（2）梯度
我们希望通过梯度下降（上升）优化策略，策略梯度可以表示为
$$\begin{gathered} \nabla \theta =\alpha \nabla J(\theta ) \\ \nabla J(\theta )=\sum _{s\in S}d(s)\sum _{a\in A}\nabla {\pi }_{\theta }(a|s)Q^{{\pi }_{\theta }}(s,a)=E_{{\pi }_{\theta }}[Q^{{\pi }_{\theta }}(s,a)\nabla \log {\pi }_{\theta }(a|s)] \end{gathered}$$

因此我们可以将$Q^{{\pi }_{\theta }}(s,a)\log {\pi }_{\theta }(a|s)$作为损失值反向传递优化模型。

1. REINFORCE

1.1 Basic

上文我们提到可以将$Q^{{\pi }_{\theta }}(s,a)\log {\pi }_{\theta }(a|s)$作为损失值反向传递优化模型。我们在强化学习（五）价值函数拟合中学到可以通过MC或TD近似Q或V函数。REINFORCE算法就是采用MC方法，利用轨迹的累计回报预估Q函数，所以策略梯度改变成为：
$$\begin{gathered} \nabla \theta =\alpha \nabla J(\theta ) \\ \nabla J(\theta )=E_{{\pi }_{\theta }}[G_t\nabla \log {\pi }_{\theta }(a|s)] \end{gathered}$$

其中G_t是一条完整轨迹获得的累积奖励。REINFORCE是一种在线学习方法，采样到的轨迹数据只能使用一次。同时因为使用累计回报$G_t$预测Q函数，所以算法的性能有一定程度的波动（高方差）。

1.2 Code

import gym
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
import rl_utils

首先定义策略网络PolicyNet，输入是某个状态，输出则是该状态下的动作概率分布。

class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=1) # dim=1，对每行使用softmax

接着定义REINFORCE算法，take_action和update是算法最重要的两个部分。在take_action中通过由PolicyNet计算获得的动作概率分布和distribution.Categorical对离散动作采样。在update中，根据与环境交互记录的轨迹计算累积回报$G_t$，并将损失函数设定为$-G_t\nabla \log {\pi }_{\theta }(a|s)$，利用梯度下降优化模型。

class REINFORCE:
    def __init__(self, state_dim, hidden_dim, action_dim, 
                learning_rate, gamma, device):
        self.policy_net = PolicyNet(state_dim, hidden_dim, 
                                    action_dim).to(device)
        self.optimizer = torch.optim.Adam(self.policy_net.parameters(),
                                         lr=learning_rate)
        self.gamma = gamma
        self.device = device
    
    # 根据动作概率分布随机采样
    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        probs = self.policy_net(state)
        
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        return action.item()
    
    def update(self, transition_dict):
        reward_list = transition_dict['rewards']
        state_list = transition_dict['states']
        action_list = transition_dict['actions']
        
        G = 0
        self.optimizer.zero_grad()
        # 通过计算累积奖励G，获得梯度
        for i in reversed(range(len(reward_list))):
            reward = reward_list[i]
            state = torch.tensor([state_list[i]],
                                dtype=torch.float).to(self.device)
            action = torch.tensor([action_list[i]]).view(-1,1).to(self.device)
            
            log_prob = torch.log(self.policy_net(state).gather(1,action))
            G = self.gamma * G + reward
            
            loss = -log_prob * G # 梯度上升
            loss.backward()
        self.optimizer.step()

问题：如何计算$\log {\pi }_{\theta }(a|s)$?怎么理解log_prob = torch.log(self.policy_net(state).gather(1,action))?

定义完策略后，我们可以在车杆环境上实验：

return_list = []
for i_episode in range(num_episodes):
    episode_return = 0
    transition_dict = {
        'states': [],
        'actions': [],
        'next_states': [],
        'rewards': [],
        'dones': []
    }
    state = env.reset()
    done = False
    while not done:
        action = agent.take_action(state)
        next_state, reward, done, _ = env.step(action)
        transition_dict['states'].append(state)
        transition_dict['actions'].append(action)
        transition_dict['next_states'].append(next_state)
        transition_dict['rewards'].append(reward)
        transition_dict['dones'].append(done)
        state = next_state
        episode_return += reward
    return_list.append(episode_return)
    # on-policy: learning from a trajectory
    agent.update(transition_dict)
    
    if (i_episode + 1) % 100 == 0:
        print(f'Episode: {i_episode + 1}, return={np.mean(return_list[-100:])}')

在1000次训练过程中我们发现，REINFORCE算法表现并不平稳，获得的return值时高时低。这是因为REINFORCE采用轨迹的累积回报预测Q函数具有较高的方差，因为采样轨迹差别过大而导致表现时好时坏。
训练完agent后，我们可以在实验环境中看看实际表现的效果如何

for i in range(5):
    episode_return = 0
    state = env.reset()
    done = False
    while not done:
        action = agent.take_action(state)
        next_state, reward, done, _ = env.step(action)
        env.render()
        state = next_state
        episode_return += reward
    print(episode_return)

---

2. Improvement Tips

正如上文所述，REINFORCE是一种具有无偏、高方差、低数据使用效率特点的On-Policy方法。我们需要设计一些技巧改进这样的方法。相关内容我在强化学习算法（五）——PPO中已经学习过了，此处尝试将他们串联起来。

2.1 Assign Suitable Credit

REINFORCE中的策略梯度为
$$\nabla J(\theta )=-G_t\nabla \log {\pi }_{\theta }(a|s)$$

其中$\nabla \log {\pi }_{\theta }(a|s)$表示轨迹中每个状态动作对发生概率修正的方向（梯度方向），$G_t$表示每个状态动作对发生概率修正的权重（大小与方向）。因此，在同一条轨迹中的所有状态动作对修正的权重相同。

很明显我们并不期望所有状态动作对的变化程度相同，所以我们需要修改G_t。我们使用状态动作对(s,a)发生后的累积奖励${\sum }_{t^{\prime }=0}^T{\gamma }^{t^{\prime }-t}{R}_{t^{\prime }}$作为每个状态动作对的修正权重，从而保证了每个状态动作对修正权重的独特性。
$$\nabla J(\theta )=-(\sum _{t^{\prime }=0}^T{\gamma }^{t^{\prime }-t}{R}_{t^{\prime }})\nabla \log {\pi }_{\theta }(a|s)$$

2.2 Add a Baseline

假设采样获得的所有状态动作对的累积奖励均大于0.在理想状态下，a, b, c三种状态动作对都被采样，状态动作的发生概率都获得了提高。但实际情况并不能保证采样到所有状态动作对，没被采样的状态动作对的发生概率不增反降。

因此我们在权重部分减去baseline，即状态在当前策略下的平均奖励$b(s_t)$。这样改变权重将不总是为正的，低于平均奖励的状态动作对发生的概率会下降，未被采样的状态动作对的概率也不会自然下降了。Basline可以使用状态价值函数衡量$b(s_t)=V^{{\pi }_{\theta }}(s_t)$。
$$\nabla J(\theta )=-(\sum _{t^{\prime }=0}^T{\gamma }^{t^{\prime }-t}{R}_{t^{\prime }}-V^{{\pi }_{\theta }}(s_t))\nabla \log {\pi }_{\theta }(a|s)$$

根据$Q^{{\pi }_{\theta }}(s,a)={\sum }_{t^{\prime }=0}^T{\gamma }^{t^{\prime }-t}{R}_{t^{\prime }}$我们可以获得新的策略梯度
$$\nabla J(\theta )=-(Q^{{\pi }_{\theta }}(s_t,a_t)-V^{{\pi }_{\theta }}(s_t))\nabla \log {\pi }_{\theta }(a_t|s_t)$$

2.3 Advantage Function

我们将$(Q^{{\pi }_{\theta }}(s_t,a_t)-V^{{\pi }_{\theta }}(s_t))$定义为优势函数$A^{\theta }(s_t,a_t)$。这样就要求我们优化两个网络：Q网络与V网络。
根据贝尔曼期望方程$Q^{{\pi }_{\theta }}(s_t,a_t)=E[R_{t+1}+\gamma V^{{\pi }_{\theta }}(s^{\prime })]$，我们可以使用单个V函数网络实现策略梯度计算。此时优势函数定义为
$$A^{{\pi }_{\theta }}(s_t,a_t)=r_t+\gamma V^{{\pi }_{\theta }}(s^{\prime })-V^{{\pi }_{\theta }}(s_t)$$

注意优势函数也就是TD-Target值了，使用优势函数表示的策略梯度为。
$$\nabla J(\theta )=-A^{{\pi }_{\theta }}(s_t,a_t)\nabla \log {\pi }_{\theta }(a_t|s_t)$$

---

3. Actor-Critic(A2C)

3.1 Basic

在介绍完策略梯度常用的三个技巧后，我们终于可以开始了解Actor-Critic框架了。后续的 TRPO、PPO、DDPG、SAC 等深度强化学习算法都是在 Actor-Critic 框架下进行发展的。深入了解 Actor-Critic 算法对读懂目前深度强化学习的研究热点有很大帮助。
采用2.1，2.2，2.3提到的技巧后，Actor网络的策略梯度表示为
$$\begin{gathered} \nabla J(\theta )=-A^{{\pi }_{\theta }}(s_t,a_t)\nabla \log {\pi }_{\theta }(a_t|s_t) \\ {A}^{{\pi }_{\theta }}(s_t,a_t)=r_t+\gamma V^{{\pi }_{\theta }}(s^{\prime })-V^{{\pi }_{\theta }}(s_t) \end{gathered}$$

A2C采用TD方法通过Actor与环境交互收集到的数据学习一个价值函数$V^{{\pi }_{\theta }}(s_t)$. Actor利用Critic学习到的价值函数优化策略网络。Critic价值网络的损失函数与梯度表示为
$$\begin{gathered} loss(w)=\frac{1}{2}(r+\gamma V_w(s_{t+1})-V_w(s_t){)}^2 \\ \nabla w=-(r+\gamma V_w(s_{t+1})-V_w(s_t))\nabla V_w(s_t) \end{gathered}$$

其中，损失函数可以理解为TD-Target与价值网络预测的均方误差。

A2C相比于REINFORCE采用MC方法直接预估Q函数，A2C性能稳定，具有较小的方差。

3.2 Code

import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
import rl_utils

首先定义策略网络PolicyNet（同REINFORCE）用于动作选择。

class PolicyNet(nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, action_dim)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=1) # dim=1，对每行使用softmax

接着定义价值网络ValueNet用于评估状态价值评估。

class ValueNet(nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super(ValueNet, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, 1)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)

Actor-Critic由actor和critic两个网络组成，critic用于输出状态的价值，actor选择动作并根据critic提供的价值优化策略。因此动作选取与REINFORCE相同，不同点在于策略优化部分。

class ActorCritic:
    def __init__(self, state_dim, hidden_dim, action_dim, actor_lr,
                critic_lr, gamma, device):
        self.actor = PolicyNet(state_dim, hidden_dim, action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                               lr = actor_lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                lr = critic_lr)
        self.gamma = gamma
        self.device = device
    
    # action selection，同REINFORCE    
    def take_action(self, state):
        state = torch.tensor([state], dtype=torch.float).to(self.device)
        probs = self.actor(state)
        action_dist = torch.distributions.Categorical(probs)
        action = action_dist.sample()
        return action.item()
    
    def update(self, transition_dict):
        states = torch.tensor(transition_dict['states'],
                             dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1,1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'],
                              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)
        
#       TD-Target = r + v(s_{t+1})
        td_target = rewards + self.gamma * self.critic(next_states) * (1-dones)
#       TD-Error
        td_error = td_target - self.critic(states)
        log_probs = torch.log(self.actor(states).gather(1, actions))
        
        actor_loss = torch.mean(-log_probs * td_error.detach())
        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
        
        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        self.actor_optimizer.step()
        self.critic_optimizer.step()

根据策略梯度与价值函数损失优化策略和价值网络。

问题：如何计算$\log {\pi }_{\theta }(a|s)$?怎么理解log_prob = torch.log(self.policy_net(state).gather(1,action))?

仍然在CartPole环境上训练测试

actor_lr = 1e-3
critic_lr = 1e-2
num_episodes = 1000
hidden_dim = 128
gamma = 0.98
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")

env_name = 'CartPole-v0'
env = gym.make(env_name)
env.seed(0)
torch.manual_seed(0)
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
agent = ActorCritic(state_dim, hidden_dim, action_dim, actor_lr, critic_lr, gamma, device)

训练部分

return_list = []
for i_episode in range(num_episodes):
    episode_return = 0
    transition_dict = {
        'states': [],
        'actions': [],
        'next_states': [],
        'rewards': [],
        'dones': []
    }
    state = env.reset()
    done = False
    while not done:
        action = agent.take_action(state)
        next_state, reward, done, _ = env.step(action)
        transition_dict['states'].append(state)
        transition_dict['actions'].append(action)
        transition_dict['next_states'].append(next_state)
        transition_dict['rewards'].append(reward)
        transition_dict['dones'].append(done)
        state = next_state
        episode_return += reward
    return_list.append(episode_return)
    # on-policy: learning from a trajectory
    agent.update(transition_dict)
    
    if (i_episode + 1) % 100 == 0:
        print(f'Episode{i_episode + 1}, return={np.mean(return_list[-100:])}')

明显的，A2C算法相对于REINFORCE的平稳性更好，训练时长缩短。
测试部分

for i in range(5):
    episode_return = 0
    state = env.reset()
    done = False
    while not done:
        action = agent.take_action(state)
        next_state, reward, done, _ = env.step(action)
        env.render()
        state = next_state
        episode_return += reward
    print(episode_return)