easy-rl/notebooks/PolicyGradient.ipynb

1. 定义算法

最基础的策略梯度算法就是REINFORCE算法，又称作Monte-Carlo Policy Gradient算法。我们策略优化的目标如下：

J_{\theta}= \Psi_{\pi} \nabla_\theta \log \pi_\theta\left(a_t \mid s_t\right)

其中$\Psi_{\pi}$在REINFORCE算法中表示衰减的回报(具体公式见伪代码)，也可以用优势来估计，也就是我们熟知的A3C算法，这个在后面包括GAE算法中都会讲到。

1.1. 策略函数设计

既然策略梯度是直接对策略函数进行梯度计算，那么策略函数如何设计呢？一般来讲有两种设计方式，一个是softmax函数，另外一个是高斯分布$\mathbb{N}\left(\phi(\mathbb{s})^{\mathbb{\pi}} \theta, \sigma^2\right)$，前者用于离散动作空间，后者多用于连续动作空间。

softmax函数可以表示为：

\pi_\theta(s, a)=\frac{e^{\phi(s, a)^{T_\theta}}}{\sum_b e^{\phi(s, b)^{T^T}}}

对应的梯度为：

\nabla_\theta \log \pi_\theta(s, a)=\phi(s, a)-\mathbb{E}_{\pi_\theta}[\phi(s,)

高斯分布对应的梯度为：

\nabla_\theta \log \pi_\theta(s, a)=\frac{\left(a-\phi(s)^T \theta\right) \phi(s)}{\sigma^2}

但是对于一些特殊的情况，例如在本次演示中动作维度=2且为离散空间，这个时候可以用伯努利分布来实现，这种方式其实是不推荐的，这里给大家做演示也是为了展现一些特殊情况，启发大家一些思考，例如Bernoulli，Binomial，Gaussian分布之间的关系。简单说来，Binomial分布，$n = 1$时就是Bernoulli分布，$n \rightarrow \infty$时就是Gaussian分布。

1.2. 模型设计

前面讲到，尽管本次演示是离散空间，但是由于动作维度等于2，此时就可以用特殊的高斯分布来表示策略函数，即伯努利分布。伯努利的分布实际上是用一个概率作为输入，然后从中采样动作，伯努利采样出来的动作只可能是0或1，就像投掷出硬币的正反面。在这种情况下，我们的策略模型就需要在MLP的基础上，将状态作为输入，将动作作为倒数第二层输出，并在最后一层增加激活函数来输出对应动作的概率。不清楚激活函数作用的同学可以再看一遍深度学习相关的知识，简单来说其作用就是增加神经网络的非线性。既然需要输出对应动作的概率，那么输出的值需要处于0-1之间，此时sigmoid函数刚好满足我们的需求，实现代码参考如下。

In [ ]:

import torch
import torch.nn as nn
import torch.nn.functional as F
class PGNet(nn.Module):
    def __init__(self, input_dim,output_dim,hidden_dim=128):
        """ 初始化q网络，为全连接网络
            input_dim: 输入的特征数即环境的状态维度
            output_dim: 输出的动作维度
        """
        super(PGNet, self).__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim) # 输入层
        self.fc2 = nn.Linear(hidden_dim,hidden_dim) # 隐藏层
        self.fc3 = nn.Linear(hidden_dim, output_dim) # 输出层
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = torch.sigmoid(self.fc3(x))
        return x

1.3. 更新函数设计

前面提到我们的优化目标也就是策略梯度算法的损失函数如下：

J_{\theta}= \Psi_{\pi} \nabla_\theta \log \pi_\theta\left(a_t \mid s_t\right)

我们需要拆开成两个部分$\Psi_{\pi}$和$\nabla_\theta \log \pi_\theta\left(a_t \mid s_t\right)$分开计算，首先看值函数部分$\Psi_{\pi}$，在REINFORCE算法中值函数是从当前时刻开始的衰减回报，如下：

G \leftarrow \sum_{k=t+1}^{T} \gamma^{k-1} r_{k}

这个实际用代码来实现的时候可能有点绕，我们可以倒过来看，在同一回合下，我们的终止时刻是$T$，那么对应的回报$G_T=\gamma^{T-1}r_T$，而对应的$G_{T-1}=\gamma^{T-2}r_{T-1}+\gamma^{T-1}r_T$，在这里代码中我们使用了一个动态规划的技巧，如下：

running_add = running_add * self.gamma + reward_pool[i] # running_add初始值为0

这个公式也是倒过来循环的，第一次的值等于：

running\_add = r_T

第二次的值则等于：

running\_add = r_T*\gamma+r_{T-1}

第三次的值等于：

running\_add = (r_T*\gamma+r_{T-1})*\gamma+r_{T-2} = r_T*\gamma^2+r_{T-1}*\gamma+r_{T-2}

In [ ]:

import torch
from torch.distributions import Bernoulli
from torch.autograd import Variable
import numpy as np

class PolicyGradient:
    
    def __init__(self, model,memory,cfg):
        self.gamma = cfg['gamma']
        self.device = torch.device(cfg['device']) 
        self.memory = memory
        self.policy_net = model.to(self.device)
        self.optimizer = torch.optim.RMSprop(self.policy_net.parameters(), lr=cfg['lr'])

    def sample_action(self,state):

        state = torch.from_numpy(state).float()
        state = Variable(state)
        probs = self.policy_net(state)
        m = Bernoulli(probs) # 伯努利分布
        action = m.sample()
        
        action = action.data.numpy().astype(int)[0] # 转为标量
        return action
    def predict_action(self,state):

        state = torch.from_numpy(state).float()
        state = Variable(state)
        probs = self.policy_net(state)
        m = Bernoulli(probs) # 伯努利分布
        action = m.sample()
        action = action.data.numpy().astype(int)[0] # 转为标量
        return action
        
    def update(self):
        state_pool,action_pool,reward_pool= self.memory.sample()
        state_pool,action_pool,reward_pool = list(state_pool),list(action_pool),list(reward_pool)
        # Discount reward
        running_add = 0
        for i in reversed(range(len(reward_pool))):
            if reward_pool[i] == 0:
                running_add = 0
            else:
                running_add = running_add * self.gamma + reward_pool[i]
                reward_pool[i] = running_add

        # Normalize reward
        reward_mean = np.mean(reward_pool)
        reward_std = np.std(reward_pool)
        for i in range(len(reward_pool)):
            reward_pool[i] = (reward_pool[i] - reward_mean) / reward_std

        # Gradient Desent
        self.optimizer.zero_grad()

        for i in range(len(reward_pool)):
            state = state_pool[i]
            action = Variable(torch.FloatTensor([action_pool[i]]))
            reward = reward_pool[i]
            state = Variable(torch.from_numpy(state).float())
            probs = self.policy_net(state)
            m = Bernoulli(probs)
            loss = -m.log_prob(action) * reward  # Negtive score function x reward
            # print(loss)
            loss.backward()
        self.optimizer.step()
        self.memory.clear()