策略梯度笔记

策略梯度

由于笔者对Actor-critic模型掌握的一直不是很好，导致最近一篇论文看的不是很明白，所以我决定好好复习一下强化学习里的Actor-critic算法，而复习Actor-critic算法，需要先从策略梯度开始，首先我们先介绍一下什么是策略梯度算法，之后我们再介绍一些实现技巧。

强化学习的组成

在强化学习的过程中，我们有三个因素，即演员，奖励函数，环境。

演员我们可以想象成一个黑箱策略网络，其由参数θ决定，输入是当前游戏的状态state，输出是一个动作。演员是根据策略来执行动作的，而策略有确定性动作和非确定性策略之分，对于非确定性策略，策略网络的输出层的神经元的个数往往与动作的个数有关，此时网络的输出为一个向量，表示采取各个动作的概率。而确定性策略则直接输出演员的动作（往往适用于连续动作空间，因为连续动作空间下动作的个数是无穷的，无法采用非确定性策略）,策略梯度为非确定性策略，即会输出一个概率向量，表示执行各动作的概率。

环境也是一个函数，输入是当前状态St和演员根据当前状态执行的动作at，输出是执行该动作后的下一个状态St+1(不是绝对的函数，对于相同的输入St和at，环境可能输出不同的St+1)

奖励函数则接受当前状态St和演员执行的动作at，给出它对演员此次动作的评分rt（此处也有说法是rt+1，不过无关紧要，表示方式而已）。

强化学习的目标

强化学习的目标就是训练演员actor，使其学习到一个策略能最大化演员在一局游戏里获得的奖励期望。

J (θ) = \sum_{τ} R (τ) p_{θ} (τ)

$J(\theta) = \sum_{\tau}R(\tau)p_{\theta}(\tau)$

其中τ为轨迹，表示一个演员从初始状态S0开始与环境进行一系列交互，直到该局游戏结束过程中遇到的状态S和演员执行的动作a的序列。

τ = {S_{0}, a_{0}, S_{1}, a_{1}, . . . . . ., S_{n}}

$\tau = \{S_0,a_0,S_1,a_1,......,S_n\}$

其中Sn表示一局游戏结束的状态。

而R(τ)表示这条轨迹带来的奖励，即r0+r1+……rn(往往需要带上折扣Γ，此处先省略，后续会补)

我们可以计算某条轨迹出现的概率：

P_{θ} (τ) = P (S_{0}) P_{θ} (a_{1} | S_{0}) P (S_{1} | S_{0}, a_{1}) . . . . . . P (S_{n} | S_{n - 1}, a_{n - 1})

$P_\theta(\tau)=P(S_0)P_\theta(a_1|S_0)P(S_1|S_0,a_1)......P(S_n|S_{n-1},a_{n-1})$

其中P为环境决定的，是我们不可控的，我们能做的就是训练Pθ，最终目的就是最大化J(θ)。将J(θ)变化一下

J (θ) = \sum_{τ} R (τ) p_{θ} (τ) = E_{τ \sim p_{θ} (τ)} (R (τ))

$J(\theta) = \sum_{\tau}R(\tau)p_{\theta}(\tau)=E_{\tau\sim p_\theta(\tau)}(R(\tau))$

最大化期望的方法

分析期望的式子，我们可以明显看出，这个式子我们能控制的只有参数θ(环境中的状态转移函数和奖励函数都是我们无法控制的)，自然能想到，是否可以通过梯度上升的方法，修改θ来增大E(θ)呢？答案是可以的。

我们先对θ求梯度：

\nabla_{θ} J (θ) = \sum_{τ} R (τ) \nabla_{θ} p_{θ} (τ)

$\nabla_\theta J(\theta)=\sum_{\tau}R(\tau)\nabla_{\theta}p_{\theta}(\tau)$

此处只有一个参数θ，所以求梯度其实和求导是一样的，我们省略梯度符号后面的θ，整理如下：

\nabla J (θ) = \sum_{τ} R (τ) \nabla p_{θ} (τ) (1)

$\nabla J(\theta)=\sum_{\tau}R(\tau)\nabla p_{\theta}(\tau) \qquad (1)$

回想一下中学我们的复合函数求导，是不是有

f^{'} (g (x)) = \frac{d f}{d g} * \frac{d g}{d x}

$f'(g(x)) =\frac{df}{dg} * \frac{dg}{dx}$

所以有

(l o g (f (x)))^{'} = \frac{1}{f (x)} * f (x)^{'}

$(log(f(x)))'=\frac{1}{f(x)} * f(x)'$

所以移项，将导数符号换成梯度符号，有

\nabla f (x) = f (x) \nabla l o g (f (x))

$\nabla f(x) = f(x) \nabla log(f(x))$

那么自然就有

\nabla p_{θ} (τ) = p_{θ} (τ) \nabla l o g (p_{θ} (τ))

$\nabla p_{\theta}(\tau)=p_{\theta}(\tau)\nabla log(p_{\theta}(\tau))$

带入（1）式：

\nabla J (θ) = \sum_{τ} R (τ) p_{θ} (τ) \nabla l o g (p_{θ} (τ)) \Rightarrow E_{τ \sim p_{θ} (τ)} (R (τ) \nabla l o g (p_{θ} (τ))

$\nabla J(\theta)=\sum_{\tau}R(\tau)p_{\theta}(\tau)\nabla log(p_{\theta}(\tau)) \\ \Rightarrow E_{\tau\sim p_\theta(\tau)}(R(\tau)\nabla log(p_{\theta}(\tau))$

进一步，我们可以把log(pθ(τ))再化简一些

\nabla l o g (p_{θ} (τ)) = \nabla l o g (P (S_{0}) P_{θ} (a_{1} | S_{0}) P (S_{1} | S_{0}, a_{1}) . . . . . . P (S_{n} | S_{n - 1}, a_{n - 1})) = \nabla (l o g (P (S_{0}) P (S_{1} | S_{0}, a_{1}) . . .) + l o g (P_{θ} (a_{1} | S_{0}) P_{θ} (a_{2} | S_{1}) . . .))

$\nabla log(p_{\theta}(\tau)) = \nabla log(P(S_0)P_\theta(a_1|S_0)P(S_1|S_0,a_1)......P(S_n|S_{n-1},a_{n-1})) \\ \qquad\qquad \qquad=\nabla (log(P(S_0)P(S_1|S_0,a_1)...)+log(P_\theta(a_1|S_0)P_\theta(a_2|S_1)...))$

两个log相加，前面的与θ无关，梯度为0，故只保留后面即可

最后可得

\nabla J (θ) = E_{τ \sim p_{θ} (τ)} (R (τ) \sum_{t = 1}^{T_{n}} l o g (P_{θ} (a_{t} | S_{t}))

$\nabla J(\theta)=E_{\tau\sim p_\theta(\tau)}(R(\tau)\sum_{t=1}^{T_n}log(P_\theta(a_t|S_t))$

由于这个期望很难求得真实值，我们往往是采样n个轨迹，然后去平均来近似这个期望：

\nabla J (θ) \approx \frac{1}{N} \sum_{n = 1}^{N} ((R (τ^{n}) \sum_{t = 1}^{T_{n}} l o g (P_{θ} (a_{t}^{n} | S_{t}^{n}))) = \frac{1}{N} \sum_{n = 1}^{N} \sum_{t = 1}^{T_{n}} R (τ^{n}) l o g (P_{θ} (a_{t}^{n} | S_{t}^{n})

$\nabla J(\theta) \approx \frac{1}{N}\sum_{n=1}^{N}((R(\tau^n)\sum_{t=1}^{T_n}log(P_\theta(a_t^n|S_t^n)))=\frac{1}{N}\sum_{n=1}^{N}\sum_{t=1}^{T_n}R(\tau^n)log(P_\theta(a_t^n|S_t^n)$

至此，我们以及有了大致的提高J(θ)的思路了，那就是通过actor与环境交互产生N条轨迹，然后根据上式求出梯度，然后根据下式更新θ。

θ = θ + η \nabla J (θ)

$\theta = \theta+\eta\nabla J(\theta)$

实现细节：

我们可以将演员的策略网络看成一个分类网络，即输入当前的状态s，输出动作预测概率向量（可以类比为分类任务中神经网络中输出层softmax后的结果）

在分类里面，我们需要最小化的是损失，也就是

J = - \sum_{i = 1}^{k} y_{i} l o g (p_{i})

$J= -\sum_{i=1}^ky_ilog(p_i)$

其中K表示类的个数，对于一个样本xi，其标签为yt，输出是一个K维向量p，当yi等于yt的适合，yi为1，否则为0

那么也就是我们需要最大化对数似然：

J = \sum_{i = 1}^{K} y_{i} l o g (p_{i})

$J = \sum_{i=1}^{K}y_{i}log(p_i)$

现在我们要最大化奖励的期望：

我们将其转换一下，在前面乘以个权重（奖励），就可以得到新的损失公式：

我们的目的就是在网络中调整参数θ以最大化该公式,实现来说，我们是每执行完一回合（即一个轨迹）就进行梯度更新，而不是采样N个轨迹，思想有些类似于随机梯度下降。

我的代码实现：

import torch
from torch import nn
from torch.nn import functional as F
import gym
import tqdm
import numpy as np
import rl_utils
import matplotlib.pyplot as plt
class Net(torch.nn.Module):
    def __init__(self , input_dim , hidden_dim , output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim , hidden_dim)
        self.fc3 = nn.Linear(hidden_dim , output_dim)

    def forward(self , x):
        x = F.relu(self.fc1(x))
        x = F.softmax(self.fc3(x) , dim = 1)
        return x

class GradientPolicy:
    def __init__(self , state_dim , hidden_dim , action_dim , lr , device):
        self.net = Net(state_dim , hidden_dim , action_dim).to(device)

        self.optimizer = torch.optim.Adam(self.net.parameters() , lr = lr)
        self.device = device

    def take_action(self , state):

        state = torch.tensor([state] , dtype = torch.float).to(self.device)
        """
        if the shape of probs is (batch_num , classes_num) , 
        the output of sample is also (batch_num , classes_num)
        """
        probs = self.net(state)

        action_dist = torch.distributions.Categorical(probs)

        action = action_dist.sample()

        return action.item()

    def learn(self , state_list , action_list , episodic_reward):
        self.optimizer.zero_grad()
        states = torch.tensor(state_list).to(self.device)
        actions = torch.tensor(action_list).to(self.device)
        action_prob = self.net(states)
        log_prob = torch.sum( torch.log(action_prob) * F.one_hot(actions , num_classes= action_prob.shape[1]))

        cost = -1.0 * log_prob * episodic_reward
        cost.backward()
        self.optimizer.step()

def runAepisode(env , agent):
    obs_list , action_list = [] , []
    res = 0
    obs , _ = env.reset()
    while True:
        obs_list.append(obs)
        action = agent.take_action(obs)
        action_list.append(action)
        obs , reward , done,_,_ = env.step(action)
        res += reward
        if done or res > 200:
            break
    return obs_list , action_list , res
if __name__ == "__main__":
    lr = 1e-3
    hidden_dim = 128
    num_episodes = 1000
    if torch.cuda.is_available():
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")
    env = gym.make("CartPole-v0")
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    agent = GradientPolicy(state_dim ,hidden_dim , action_dim , lr , device)
    return_list = []
    for i in range(num_episodes):
        states , actions , rewards = runAepisode(env , agent)
        print(i , rewards)
        return_list.append(rewards)

        agent.learn(states , actions , rewards)

    torch.save(agent.net.state_dict() , 'model.pth')
    episodes_list = list(range(len(return_list)))
    plt.plot(episodes_list, return_list)
    plt.xlabel('Episodes')
    plt.ylabel('Returns')
    plt.show()