【强化学习】策略梯度算法中的损失函数

策略梯度算法推导

适合有一定强化学习算法基础的人看。

首先我们定义一条轨迹的奖励，有两种定义方法：
无折扣有限步奖励
$$R(\tau )=\sum _{t=0}^T{r}_t.$$
折扣无限步奖励
$$R(\tau )=\sum _{t=0}^{\infty }{\gamma }^t{r}_t.$$

以上两种轨迹奖励的求和方法第一种比较简单，第二种比较合适。而且一般对于第二种也是有限的，比如游戏必然会终结的。

智能体和环境交互的过程中，产生的轨迹有很多条，强化学习的目标是使得奖励的期望值最大：
假设在环境中智能体以策略$\pi$和环境进行交互产生一条轨迹$\tau$的概率：
$$P(\tau |\pi )={\rho }_0(s_0)\prod _{t=0}^{T-1}P(s_{t+1}|s_t,a_t)\pi (a_t|s_t).$$
其中${\rho }_0(s_0)$是初始状态出现的概率，$P(s_{t+1}|s_t,a_t)$是状态转移的概率，$\pi (a_t|s_t)$是策略，也就是在状态$s_t$下采取动作$a_t$的概率。

那么目标就是最大化期望奖励：
$$J(\pi )={\int }_{\tau }P(\tau |\pi )R(\tau )=\underset{\tau \sim \pi }{E}[R(\tau )].$$
强化学习的核心问题是找到一个策略，使得目标函数值最大：
$${\pi }^{\ast }=\arg \underset{\pi }{\max }J(\pi )$$

目标函数的自变量是$\pi$，如果我们用某种方法来表示策略，参数是$\theta$，那么上面的式子就可以写成：
$$J({\pi }_{\theta })={\int }_{\tau }P(\tau |{\pi }_{\theta })R(\tau )=\underset{\tau \sim {\pi }_{\theta }}{E}[R(\tau )].$$
最终目标函数是和参数$\theta$有关的，我们只需要更新参数$\theta$就能影响策略$\pi (a|s)$的分布，进而影响$J({\pi }_{\theta })$.

接下来我们求目标函数的梯度
1.轨迹$\tau$出现的概率
$$P(\tau |\theta )={\rho }_0(s_0)\prod _{t=0}^{T}P(s_{t+1}|s_t,a_t){\pi }_{\theta }(a_t|s_t).$$
2.对轨迹的概率$P(\tau |\theta )$求关于$\theta$的导数
这里用到了一个技巧后面会用到
$${\nabla }_{\theta }P(\tau |\theta )=P(\tau |\theta ){\nabla }_{\theta }\log P(\tau |\theta ).$$
$\log P(\tau |\theta )$可以展开为
$$\log P(\tau |\theta )=\log {\rho }_0(s_0)+\sum _{t=0}^T(\log P(s_{t+1}|s_t,a_t)+\log {\pi }_{\theta }(a_t|s_t)).$$
第一项和第二项都和参数无关，只有最后一项和参数有关：
$${\nabla }_{\theta }logP(\tau |\theta )=\sum _{t=0}^T{\nabla }_{\theta }log{\pi }_{\theta }(a_t|s_t)$$
所以
$$\begin{array}{rlr}{\nabla }_{\theta }J({\pi }_{\theta })& ={\nabla }_{\theta }\underset{\tau \sim {\pi }_{\theta }}{E}[R(\tau )]& \\ & ={\nabla }_{\theta }{\int }_{\tau }P(\tau |\theta )R(\tau )& \text{展开期望}\\ & ={\int }_{\tau }{\nabla }_{\theta }P(\tau |\theta )R(\tau )& \text{把梯度放到里面}\\ & ={\int }_{\tau }P(\tau |\theta ){\nabla }_{\theta }\log P(\tau |\theta )R(\tau )& \text{log技巧}\\ & =\underset{\tau \sim {\pi }_{\theta }}{E}{\nabla }_{\theta }\log P(\tau |\theta )R(\tau )& \text{又重新变成期望的格式}\\ \therefore {\nabla }_{\theta }J({\pi }_{\theta })& =\underset{\tau \sim {\pi }_{\theta }}{E}\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )& \text{目标函数的梯度}\end{array}$$

最后的表达式就是我们所需要的目标函数的梯度，后面的$R(\tau )$部分还可以写成其他的形式，这里不做讨论，例如写成Rewart-to-go的形式：
$${\nabla }_{\theta }J({\pi }_{\theta })=\underset{\tau \sim {\pi }_{\theta }}{E}\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\sum _{t^{\prime }=t}^{T}R(s_{t^{\prime }},a_{t^{\prime }},s_{t^{\prime }+1}).$$
还可以减去一个baseline来减小方差
$${\nabla }_{\theta }J({\pi }_{\theta })=\underset{\tau \sim {\pi }_{\theta }}{E}\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)\left(\sum _{t^{\prime }=t}^{T}R(s_{t^{\prime }},a_{t^{\prime }},s_{t^{\prime }+1})-b(s_t)\right).$$
这个式子是正确的，因为后面的baseline不会影响梯度的值，可以证明第二项的期望为0.

回到话题，我们已经求出了目标函数的梯度：
$${\nabla }_{\theta }J({\pi }_{\theta })=\underset{\tau \sim {\pi }_{\theta }}{E}\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau )$$
最外层的期望可以通过蒙特卡洛采样的方法来实现，即采样多条轨迹然后求平均值：
$$\hat{g}=\frac{1}{|\mathcal{D}|}\sum _{\tau \in \mathcal{D}}\sum _{t=0}^T{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)R(\tau ),$$

$\mathcal{D}$是采集的轨迹的集合：$\mathcal{D}=\{{\tau }_1,{\tau }_2,...,{\tau }_{|\mathcal{D}|}\}$

但是问题是，放到神经网络中，参数该如何更新呢？

策略梯度算法的损失函数

神经网络的结构

如果我们用神经网络来表示策略${\pi }_{\theta }$,神经网络的输入就是状态$s_t$，输出是每一个动作的概率${\pi }_{\theta }(a_t|s_t)$，$\theta$是神经网络中的参数.

这里考虑的是离散动作的情况，如果是连续动作，输出是高斯分布的参数。如果是多维连续动作，那么输出是多维高斯分布的参数。

要想输出是概率，那么最后一层之后我们需要接一个$softmax$函数，这样才能保证每一个输出值大于0，且和为1.
$$softmax(x_i)=\frac{e^{x_i}}{{\sum }_{j=1}^N{e}^{x_j}}$$

Note：神经网络的参数$\theta$影响的是策略${\pi }_{\theta }$的分布，策略的分布，影响的是智能体在环境中的交互产生的轨迹的分布，轨迹的分布影响的是最终智能体得到的累计回报。

我们希望产生回报大的轨迹出现的概率高，产生回报小的轨迹出现的概率低。如果能够构造一个好的动作评判指标，来判断一个动作的好和不好，那么就可以通过改变动作出现的概率来优化策略。
我们构造损失函数如下：
$$L(\theta )=-\sum _{\tau \sim {\pi }_{\theta }}log{\pi }_{\theta }(a|s)R(\tau )$$
那么在深度学习框架中实现参数更新只需要如下的伪代码：

loss = L(theta)
loss.backward() #梯度反向传播
loss.step() #参数theta更新

上述代码会对loss求梯度，然后通过梯度下降来更新参数。

回顾一下交叉熵：
$$H(p,q)=-\sum _{x}p(x)logq(x)$$
其中p和q是两个概率分布，交叉熵刻画的是两个概率分布之间的距离，两个概率的分布越接近，交叉熵越小。
回到我们的话题，我们希望当轨迹累计回报$R(\tau )$ 越大的时候，对应的动作出现的概率也越大 $log{\pi }_{\theta }(a_t|s_t)$. 反之亦然。 换句话说，我们希望$R(\tau )$和$log{\pi }_{\theta }(a_t|s_t)$的分布越接近越好，用交叉熵损失函数是再合适不过了！

这里还可以继续深入理解： 使用交叉熵损失函数将两者绑定在了一起，其中$R(\tau )$在确定环境中给定策略下是确定的，无法改变的。智能体在和环境交互的过程中，发现轨迹${\tau }_i$的奖励大一些，于是就通过交叉熵损失函数将对应的动作出现的概率也增大了一些（通过更新参数实现）。于是，智能体再次探索环境的时候，走那些奖励大一些的轨迹的概率更大一些（这就是利用已经学会的知识）。

策略梯度算法的简单实现

声明：以下代码为大概思路，不是完整的代码。完整代码可以参考 【强化学习】spinningup最简单的策略梯度（VPG）代码详细注释——基于pytorch实现
或者这篇：动手强化学习（九）：策略梯度算法 代码通俗易懂，适合新手。

定义神经网络

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)

神经网络就两层，第一层的输出接一个激活函数，第二层的输出接一个softmax函数，目的是得到输入状态为$s_t$时，输出动作的概率。

选取动作

神经网络的输出是动作的概率分布，根据这个分布选取当前智能体执行的动作是什么：

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()
        # 返回动作(这里的动作是index，从0开始)
        return action.item()

采样

接下来就是智能体和环境进行交互采样，得到一个批次的数据。

损失函数

定义损失函数

    def compute_loss(obs, act, weights):
        # 注意：这里传入的是批次数据（多条轨迹），所以最后求平均值，为平均loss，或者
        logp = get_policy(obs).log_prob(act) # 计算 logp(a|s)
        return -(logp * weights).mean() # 交叉熵损失函数
        

参数更新

有了损失函数之后，定义优化器，更新参数：

optimizer = Adam(logits_net.parameters(), lr=lr)
# 采样之后，计算梯度，更新网络参数
optimizer.zero_grad()
# 这里是计算梯度
batch_loss = compute_loss(obs=torch.as_tensor(batch_obs, dtype=torch.float32),
                                  act=torch.as_tensor(batch_acts, dtype=torch.int32),
                                  weights=torch.as_tensor(batch_weights, dtype=torch.float32)
# 梯度反向传播
batch_loss.backward()
optimizer.step()