DQN实战CartPole

这篇博文要讲解的是利用DQN来做CartPole任务

---

回报的定义：

我们知道，给定一个状态$s$，agent根据策略$\pi (a|s)$做出行为$a$，得到的奖励是$r$，然后环境根据状态转移概率$P(s^{\prime }|s)$转移到新的状态$s^{\prime }$.

强化学习中更多时候关注的是给定某一个状态$S_t$，它的累计奖励，也叫“回报”或者“收获”，用英文return表示。

定义给定状态$S_t$下的回报：
$$G_t=R_t+\gamma R_{t+1}+{\gamma }^2{R}_{t+2}+\cdots$$

它的含义就是在状态$S_t$下开始，根据策略做出行为，获得奖励，累计奖励，更新状态，直到终止状态下所有奖励的累计。加入折扣因子$\gamma$是因为未来的奖励对当前的状态具有不确定性，避免在计算回报时陷入循环。

---

动作价值函数(action-value function)

我们可以看出，一个状态的回报反映了该状态在一个状态序列下的重要程度。但是当给定一个状态时，它可能产生无数多个状态序列。
具体的：

• 一方面，给定状态$s$下，根据策略函数$\pi (a|s)$可能会得到不同的行为，这就有了不同的奖励，显然回报不同。
• 另一方面，给定状态$s$下，根据状态转移矩阵$P(s^{\prime }|s)$，同样是根据概率随机的转移到下一个状态$s^{\prime }$，显然最终的回报也是不同的。

价值函数(value function):

所以简单地用回报$G_t$是不能准确的定义一个状态的好坏的，因为$G_t$仅仅针对某一个状态序列。

那么一种自然的想法就是状态$s$下，根据策略$\pi (a|s)$和状态转移矩阵$P(s^{\prime }|s)$，穷举所有可能产生的状态序列下的回报${\sum }_{t=1}^{\infty }{G}_t$，作为一个指标，衡量状态$s$的好坏.

这也就是对给定状态$s$下所有可能产生的状态序列下的回报的期望。
$$\mathbb{E}[G_t|S_t=s]$$

这就是强化学习中“价值”的概念，也就是状态$s$的价值就是$v_{\pi }(s)=\mathbb{E}[G_t|S_t=s]$. value的含义就是状态回报的期望.

如果存在一个函数，给定一个状态，就能得到该状态对应的价值，那么这个函数就叫做价值函数(value function).
价值函数是状态到价值的映射$state\stackrel{valuefunction}{}value$

注意：价值函数是基于策略$\pi$的。

最优价值函数：

我们知道在不同的策略$\pi$下，给定状态$s$，$v_{\pi }(s)$显然是不同的。
最优价值函数就是对所有的策略$\pi$下，$v_{\pi }(s)$值最大的价值函数：
$$v_{\ast }(s)=\underset{\pi }{\max }{v}_{\pi }(s)$$

动作价值函数:

价值函数评估的是根据策略$\pi$当前状态的好坏，动作价值函数评估的是当前状态下根据策略$\pi$做出某一个行为的好坏。所以它引入了动作(也就是行为)的概念。
$$q_{\pi }(s,a)=\mathbb{E}[G_t|S_t=s,A_t=a]$$

再啰嗦一遍:

• $v_{\pi }(s)=\mathbb{E}[G_t|S_t=s]$的含义是基于策略$\pi$，给定状态$s$下，所能得到的回报的期望，评估的是状态$s$的好坏。
• $q_{\pi }(s,a)=\mathbb{E}[G_t|S_t=s,A_t=a]$的含义是基于策略$\pi$，给定状态$s$下做出行为$a$之后，所能得到的回报的期望，评估的是状态$s$下做出行为$a$的好坏。

最优动作价值函数(optimal action-value function)

我们知道在不同的策略$\pi$下，给定状态$s$，$q_{\pi }(s,a)$显然是不同的。
最优动作价值函数就是对所有的策略$\pi$下，$q_{\pi }(s,a)$值最大时对应的动作价值函数:

$$q_{\ast }(s,a)=\underset{\pi }{\max }{q}_{\pi }(s,a)$$

---

TD学习(Temporal Difference时序差分)：

我们来推导动作价值函数的Bellman equation(贝尔曼等式)：

$$\begin{array}{rl}{q}_{\pi }(s,a)& =\mathbb{E}[G_t|S_t=s,A_t=a]\\ & =\mathbb{E}[R_t+\gamma R_{t+1}+{\gamma }^2{R}_{t+2}+\cdots |S_t=s,A_t=a]\\ & =\mathbb{E}[R_t+\gamma (R_{t+1}+\gamma R_{t+2}+{\gamma }^2{R}_{t+3}+\cdots )|S_t=s,A_t=a]\\ & =\mathbb{E}[R_t+\gamma G_{t+1}|S_t=s,A_t=a]\\ & =\mathbb{E}[R_t|S_t=s,A_t=a]+\gamma \mathbb{E}[G_{t+1}|S_t=s,A_t=a]\\ & =\mathbb{E}[R_t|S_t=s,A_t=a]+\gamma q_{\pi }(S_{t+1},A_{t+1})\\ & =\mathbb{E}[R_t|S_t=s,A_t=a]+\gamma \mathbb{E}[q_{\pi }(S_{t+1},A_{t+1})|S_t=s,A_t=a]\\ & =\mathbb{E}[R_t+\gamma q_{\pi }(S_{t+1},A_{t+1})|S_t=s,A_t=a]\end{array}$$

整个推导过程中需要注意的几点：

• $G_{t+1}=R_{t+1}+\gamma R_{t+2}+{\gamma }^2{R}_{t+3}+\cdots$
• $q_{\pi }(S_{t+1},A_{t+1})=\mathbb{E}[G_{t+1}|S_t=s,A_t=a]$
• 由于$q_{\pi }(S_{t+1},A_{t+1})$的输出就是一个数值，又因为$\mathbb{E}(a)=a$，也就是说对常数的期望等于该常数。所以$q_{\pi }(S_{t+1},A_{t+1})=\mathbb{E}[q_{\pi }(S_{t+1},A_{t+1})]$

上面的公式中由于有期望，所以我们一般是通过采样的方法近似这个期望值，也就是用观测得到的奖励$r_t$近似$R_t$，观测得到的下一时刻的状态$s_{t+1},a_{t+1}$代替$S_{t+1},A_{t+1}$：
$$q_{\pi }(s,a)\approx r_t+\gamma q_{\pi }(s_{t+1},a_{t+1})$$
所以不难得出价值函数的Bellman equation就是：
$$v_{\pi }(s)=\mathbb{E}[R_t+\gamma v_{\pi }(S_{t+1})|S_t=s]$$

近似值就是：
$$v_{\pi }(s)\approx r_t+\gamma v_{\pi }(s_{t+1})$$

前面我们已经提到，价值函数评估的是一个状态的价值，动作价值函数评估的是该状态下做出某一个动作的价值。那么怎么评估这个价值呢？

• 一种直观的想法是在多个状态序列中状态收获的平均值可以近似为价值。这种做法就是MC学习(蒙特卡罗学习)，它的计算方式是计算一个序列中所有奖励的总和作为这个状态序列中初始状态的收获。然后累计多个状态序列的收获的平均值作为价值。
• 另一种就是时序差分学习方法，它不需要直到序列结束才能评估价值。具体的就是：在估算某一个状态的价值时，利用离开该状态的即时奖励$r_t$和下一时刻状态$s_{t+1}$的预估状态价值乘以衰减系数$\gamma$。公式如下:

$$v_{\pi }(s_t)\leftarrow v_{\pi }(s_t)-[\alpha (r_t+\gamma v_{\pi }(s_{t+1})-v_{\pi }(s_t))]$$
其中的$r_t+\gamma v_{\pi }(s_{t+1})$称为TD target，$r_t+\gamma v_{\pi }(s_{t+1})-v_{\pi }(s_t)$称为TD error。

TD学习有两种算法：SARSA和Q-learning.

• SARSA算法用来学习价值函数或者动作价值函数。
• Q-learning算法用来学习最优价值函数或者最优动作价值函数。

我们主要来看如何用两种算法来学习动作价值函数和最优动作价值函数，按照时序差分的思想，我们可以推导出：

SARSA算法：

$$q_{\pi }(s_t,a_t)\leftarrow q_{\pi }(s_t,a_t)-[\alpha (r_t+\gamma q_{\pi }(s_{t+1},a_{t+1})-q_{\pi }(s_t,a_t))]$$

Q-learning算法：

前面我们已经推导出动作价值函数的Bellman equation是:
$$q_{\pi }(s,a)=\mathbb{E}[R_t+\gamma q_{\pi }(S_{t+1},A_{t+1})|S_t=s,A_t=a]$$
最优动作价值函数定义为：
$$q_{\ast }(s,a)=\underset{\pi }{\max }{q}_{\pi }(s,a)$$
所以最优动作价值函数的Bellman equation为：
$$q_{\ast }(s,a)=\mathbb{E}[R_t+\gamma q_{\ast }(S_{t+1},A_{t+1})|S_t=s,A_t=a]$$

$q_{\ast }(s,a)$可评估给定状态下所有动作的好坏，那么当状态到达$S_{t+1}$时，最优动作自然是$A_{t+1}=\underset{a^{\prime }\in A}{\mathrm{argmax}}{q}_{\ast }(S_{t+1},a^{\prime })$
所以我们可以推出来：
$$q_{\ast }(S_{t+1},A_{t+1})=\underset{a^{\prime }\in A}{\max }{q}_{\ast }(S_{t+1},a^{\prime })$$
所以最优动作价值函数的Bellman equation还可以写成：
$$q_{\ast }(s,a)=\mathbb{E}[R_t+\gamma \underset{a^{\prime }\in A}{\max }{q}_{\ast }(S_{t+1},a^{\prime })|S_t=s,A_t=a]$$

仍然用采样的方式去掉期望，也就是用观测得到的奖励$r_t$，以及观测得到的$s_{t+1}$代替上式的$S_{t+1}$。
所以我们得到最优动作价值函数$q_{\ast }(s,a)$的近似值：

$$q_{\ast }(s,a)\approx r_t+\gamma \underset{a^{\prime }\in A}{\max }{q}_{\ast }(s_{t+1},a^{\prime })$$

Q-learning算法的更新过程就是：
$$q_{\ast }(s,a)\leftarrow q_{\ast }(s,a)-\alpha (r_t+\gamma \underset{a^{\prime }\in A}{\max }{q}_{\ast }(s_{t+1},a^{\prime })-q_{\ast }(s,a))$$
其中$r_t+\gamma \underset{a^{\prime }\in A}{\max }{q}_{\ast }(s_{t+1},a^{\prime })$称为TD target，$r_t+\gamma \underset{a^{\prime }\in A}{\max }{q}_{\ast }(s_{t+1},a^{\prime })-q_{\ast }(s,a)$称为TD error。

---

DQN

DQN就是利用神经网络近似最优动作价值函数$q_{\ast }(s,a)$。我们写成$q(s,a;\theta )$,$\theta$就是网络的参数。我们的目的现在变为学习参数$\theta$，使得其更好的近似最优动作价值函数。

• 神经网络的输入是state
• 神经网络的输出是对每一个action的评分
• 我们选取分数最高的动作(argmax)，然后执行这个动作

定义神经网络的损失函数是均方误差。也就是
$$loss=\frac{1}{2}\sum _t(q(s_t,a_t;\theta )-y_t{)}^2$$
算法流程：

1. 收集一次transition: ($s_t,a_t,r_t,s_{t+1}$)
2. 计算TD target: $y_t=r_t+\gamma \underset{a^{\prime }\in A}{\max }q(s_t,a^{\prime };\theta )$
3. 计算TD error: ${\delta }_t=q(s_t,a_t;\theta )-y_t$
4. 计算梯度为${\delta }_t\ast \frac{\partial q(s_t,a_t;\theta )}{\partial \theta }$
5. 更新参数$\theta \leftarrow \theta -\alpha \ast {\delta }_t\ast \frac{\partial q(s_t,a_t;\theta )}{\partial \theta }$

此外DQN的训练过程有许多的技巧，比如：

• 经验回放
• $ϵ$贪婪策略
• DDQN

经验回放(experience replay)

我们把一个$(s_t,a_t,r_t,s_{t+1})$称为一个transition.按照上面的算法流程我们可以看到，我们是拿一个transition，去更新一次参数$\theta$，然后这个transition就丢弃不用了。然后根据新的transition$(s_{t+1},a_{t+1},r_{t+1},s_{t+2})$再重复这个过程。

我们把从初始状态$s_0$到终止状态$s_n$之间的所有transition称为经验experience

这种做法有两点问题：

1. 浪费了很多的transition，这些transition是可以重复利用的。(这些transition其实就是数据，大家在做深度学习时，就算batch_size设置为1，但是每一个样本也经过了很多次epoch，也就是重复利用了epoch次。而这里相当于用一次就丢了。)
2. 如果按照上述的算法流程，我们是序列式的将各个transition送到模型中，这就导致各个transition之间有较强的关联性。(深度学习中，我们的训练样本都是独立同分布的，各个样本之间毫无关联。)

所以就有了经验回放这种做法：

有一个replay buffer池，收集很多的transition, $(s_t,a_t,r_t,s_{t+1}),(s_{t+1},a_{t+1},r_{t+1},s_{t+2}),(s_{t+2},a_{t+2},r_{t+2},s_{t+2}),\cdots ,(s_{t+n},a_{t+n},r_{t+n},s_{t+n+1})$,这里的n就是池的容量，将这些transition放到replay buffer中，注意这里我是顺序写的，但实际上不是的，各个transition之间是打乱了顺序的。然后从replay buffer中取出batch_size个transition更新参数。

$ϵ$贪婪策略

在模型刚开始训练的时候，网络的参数都是随机初始化的，所以不会很好的近似最优价值函数。所以在刚开始我们有一定的概率随机做出一个行为。在训练了很多步骤后，网络已经能很好的近似最优价值函数，那么我们此时倾向于做出神经网络输出的行为，而不再是随机做出行为。

$ϵ$-greedy策略是探索(exploration)和利用(exploitation)之间的折中。

if random.random()>epsilon_threshold:
	return Q_network(state)
else:
	return random.randint(a=0,b=action_space-1)#随机做出行为

DDQN

Double DQN是解决DQN高估问题的一个非常有效的方法。

$$q_{\pi }(s_t,a_t)=r_t+\gamma \underset{a^{\prime }}{\max }{q}_{\pi }(s_{t+1},a^{\prime })$$

上式中有$\underset{a^{\prime }}{\max }{q}_{\pi }(s_{t+1},a^{\prime })$这一项，这一项是对状态$s_{t+1}$所有可能的动作价值的估计中最大的，这就会造成高估问题。举个例子：

$x_1,x_2,\cdots ,x_n$是n个数据点，我们向其中添加均值为0的噪声，此时的数据表示为${\delta }_1,{\delta }_2,\cdots ,{\delta }_n$
那么我们可以证明如下式成立：
$$\begin{gathered} \underset{i}{\mathrm{mean}}(x_i)=\underset{i}{\mathrm{mean}}({\delta }_i) \\ \underset{i}{\max }(x_i)\le \underset{i}{\max }({\delta }_i) \end{gathered}$$
这是因为噪声的均值是0，所以添加噪声后的数据的均值不变，但是既然均值为0，那么说明噪声中肯定既有负值又有正值，所以导致

• 添加了噪声后的数据的最大值变大
• 添加了噪声后的数据的最小值变小

我们回到DQN，$\underset{a^{\prime }}{\max }{q}_{\pi }(s_{t+1},a^{\prime })$这一项就是对所有的动作进行估计，估计它的价值。我们可以将$Q$函数中的参数视为均值是0的噪声，那么我们就知道了$\underset{a^{\prime }}{\max }{q}_{\pi }(s_{t+1},a^{\prime })$就会导致估计值大于真实值。从而使得$q_{\pi }(s_t,a_t)=r_t+\underset{a^{\prime }}{\max }{q}_{\pi }(s_{t+1},a^{\prime })$变大。

我们从上面的公式看到，DQN在估计$t$时刻的$Q$值的时候，利用了$t+1$时刻$Q$值的估计，这就是自举(booststrapping)，我们已经知道$t+1$时刻由于最大化操作造成了高估，那么就会使得DQN对$t$时刻的估计同样变大，造成高估问题。以此下去，使得高估问题越来越严重。

通过上面的理论说明我们知道造成DQN高估的问题有两个，一个是最大化；另一个是自举。

Double DQN就是有两个Q-Network，记为$Q$和$Q^{\prime }$，

• $Q$用来收集轨迹，选出action $a^{\ast }=\underset{a^{\prime }\in A}{\mathrm{argmax}}Q(s_{t+1},a^{\prime })$
• $Q^{\prime }$根据$Q$选出来具有最大价值的行为计算target，$r_t+\gamma Q^{\prime }(s_{t+1},a^{\ast })$

我们对比一下DQN和DDQN:

• DQN: $Q(s_t,a_t)=r_t+\gamma \underset{a^{\prime }\in A}{\max }Q(s_{t+1},a^{\prime })$
• DDQN: $Q(s_t,a_t)=r_t+\gamma Q^{\prime }(s_{t+1},\underset{a^{\prime }\in A}{\mathrm{argmax}}Q(s_{t+1},a^{\prime }))$

可以看出DDQN减轻了最大化和自举带来的高估问题。

• 因为即使$\underset{a^{\prime }\in A}{\mathrm{argmax}}Q(s_{t+1},a^{\prime })$出现高估，但是$Q^{\prime }$未必会对高估出来的动作$a^{\ast }$高估，这就减轻了最大化导致的高估
• 由于我们是用$Q^{\prime }$来估计$t+1$时刻动作的价值，然后作为target更新$Q$在$t$时刻的估计，这就减轻了自举导致的高估问题。

其中$Q$和$Q^{\prime }$的网络结构完全相同，对于$Q$我们是实时更新的。对于$Q^{\prime }$，我们每隔一定时间步，将$Q$的参数复制给$Q^{\prime }$，用pytorch实现就是：

Q'.load_state_dict(Q.state_sict())

代码部分

在进行代码之前，我们写一下伪代码再熟悉一遍流程：

env=gym.make("CartPole-v")#创建环境
while True:
	state=env.reset()#每一个序列episode结束后重置环境，也就是让小车回到原点
	for t in range(T):
		action=Q_network(state)#我们利用神经网络作为Q_network近似最优价值函数，它的输出就是最优动作
		next_state,reward,done,info=env.step(action)#我们执行这个动作，获得了奖励，并且更新到下一个状态
		#done表示的就是是否达到了终止状态,info不用考虑
		if done:
			break

必要说明

在正常的CartPole游戏中，状态是由四个数值组成的，分别代表(车的位置，车的加速度，杆的角度，杆的角加速度)。
但是我们这里有另一种做法，就是把当前画面和前一帧画面之间的差值作为状态输入给神经网络。 神经网络的输出是对行为的打分，这里只有两个行为(左和右)

导入必要的包

import gym
import math
import random
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple
from itertools import count
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T
env = gym.make('CartPole-v0').unwrapped#创建环境
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display
plt.ion()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

构造状态

resize = T.Compose([T.ToPILImage(),
                    T.Resize(40, interpolation=Image.CUBIC),
                    T.ToTensor()])

# This is based on the code from gym.
screen_width = 600


def get_cart_location():
    world_width = env.x_threshold * 2
    scale = screen_width / world_width
    return int(env.state[0] * scale + screen_width / 2.0)  # MIDDLE OF CART


def get_screen():
    screen = env.render(mode='rgb_array').transpose(
        (2, 0, 1))  # transpose into torch order (CHW)
    # Strip off the top and bottom of the screen
    screen = screen[:, 160:320]
    view_width = 320
    cart_location = get_cart_location()
    
    if cart_location < view_width // 2:
        slice_range = slice(view_width)
    elif cart_location > (screen_width - view_width // 2):
        slice_range = slice(-view_width, None)
    else:
        slice_range = slice(cart_location - view_width // 2,
                            cart_location + view_width // 2)
    # Strip off the edges, so that we have a square image centered on a cart
    screen = screen[:, :, slice_range]
    # Convert to float, rescare, convert to torch tensor
    # (this doesn't require a copy)
    screen = np.ascontiguousarray(screen, dtype=np.float32) / 255
    screen = torch.from_numpy(screen)
    # Resize, and add a batch dimension (BCHW)
    return resize(screen).unsqueeze(0).to(device)


env.reset()
plt.figure()
plt.imshow(get_screen().cpu().squeeze(0).permute(1, 2, 0).numpy(),
           interpolation='none')
plt.title('Example extracted screen')
plt.show()

也就是说调用get_screen()函数会返回如下图片

我们将状态定义为两个相邻图片之间的差值，具体的：

last_screen=get_screen()
current_screen=get_screen()
state=current_screen-last_screen
print(state.size())

我们可以得到状态的形状是(1,3,40,80)。1代表的一个样本，3是通道数，40是高，80是宽。

构造Deep Q Network

我们定义一个由三层卷积构成的神经网络近似最优动作价值函数。我们把这个网络称为Q network.
它的输入就是状态，形状是(batch_size,3,40,80)
输出就是给每一个状态对应的动作打的分数，形状是(batch_size,2)

class DQN(nn.Module):

    def __init__(self):
        super(DQN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=5, stride=2)
        self.bn1 = nn.BatchNorm2d(16)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=2)
        self.bn2 = nn.BatchNorm2d(32)
        self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2)
        self.bn3 = nn.BatchNorm2d(32)
        self.head = nn.Linear(448, 2)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        return self.head(x.view(x.size(0), -1))

构建replay_buffer

Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward'))


class ReplayMemory(object):

    def __init__(self, capacity):
        self.capacity = capacity
        self.memory = []
        self.position = 0

    def push(self, *args):
        """Saves a transition."""
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.position] = Transition(*args)
        self.position = (self.position + 1) % self.capacity

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size)

    def __len__(self):
        return len(self.memory)

利用$ϵ$-greedy策略选择行为

选择epsilon贪婪策略，一开始时尽可能的选取随机策略，因为此时的网络还不能很好的近似最优价值函数。所以此时我们需要尽可能的探索每一个状态下可能出现的每一个行为。

当一定步骤之后，我们的网络已经可以给出最优动作的时候，我们倾向于使用这个最优动作，而不再随机做出行为。

$ϵ$-greedy策略的目的就是再探索和利用中取得平衡

BATCH_SIZE = 128
GAMMA = 0.999
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 200
TARGET_UPDATE = 10

policy_net = DQN().to(device)
target_net = DQN().to(device)
target_net.load_state_dict(policy_net.state_dict())
target_net.eval()

optimizer = optim.RMSprop(policy_net.parameters())
memory = ReplayMemory(10000)


steps_done = 0


def select_action(state):
    global steps_done
    sample = random.random()
    eps_threshold = EPS_END + (EPS_START - EPS_END) * \
        math.exp(-1. * steps_done / EPS_DECAY)
    steps_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            return policy_net(state).max(1)[1].view(1, 1)
    else:
        return torch.tensor([[random.randrange(2)]], device=device, dtype=torch.long)


episode_durations = []


def plot_durations():
    plt.figure(2)
    plt.clf()
    durations_t = torch.tensor(episode_durations, dtype=torch.float)
    plt.title('Training...')
    plt.xlabel('Episode')
    plt.ylabel('Duration')
    plt.plot(durations_t.numpy())
    # Take 100 episode averages and plot them too
    if len(durations_t) >= 100:
        means = durations_t.unfold(0, 100, 1).mean(1).view(-1)
        means = torch.cat((torch.zeros(99), means))
        plt.plot(means.numpy())

    plt.pause(0.001)  # pause a bit so that plots are updated
    if is_ipython:
        display.clear_output(wait=True)
        display.display(plt.gcf())

def optimize_model():
    if len(memory) < BATCH_SIZE:
        return#当存储区中的transition不足batch_size的时候，我们暂时不更新参数
    transitions = memory.sample(BATCH_SIZE)#采样出batch_size个transition
    batch = Transition(*zip(*transitions))

    #需要注意的是，当某个状态下做出行为后，很有可能就导致游戏结束了，所以next_state此时应该是None，对应的done是1
    non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
                                          batch.next_state)), device=device, dtype=torch.bool)
    non_final_next_states = torch.cat([s for s in batch.next_state
                                                if s is not None])
    state_batch = torch.cat(batch.state)
    action_batch = torch.cat(batch.action)
    reward_batch = torch.cat(batch.reward)

    # Compute Q(s_t, a) - the model computes Q(s_t), then we select the
    # columns of actions taken
    state_action_values = policy_net(state_batch).gather(1, action_batch)

    # Compute V(s_{t+1}) for all next states.
    next_state_values = torch.zeros(BATCH_SIZE, device=device)
    next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0].detach()
    #对于那些next_state是游戏结束的state，我们将它们的next_state_values设置为0
    # Compute the expected Q values
    expected_state_action_values = (next_state_values * GAMMA) + reward_batch

    # Compute Huber loss
    loss = F.smooth_l1_loss(state_action_values, expected_state_action_values.unsqueeze(1))

    # Optimize the model
    optimizer.zero_grad()
    loss.backward()
    for param in policy_net.parameters():
        param.grad.data.clamp_(-1, 1)
    optimizer.step()

num_episodes = 5000
for i_episode in range(num_episodes):
    # Initialize the environment and state
    env.reset()
    last_screen = get_screen()
    current_screen = get_screen()
    state = current_screen - last_screen
    for t in count():
        # Select and perform an action
        action = select_action(state)
        _, reward, done, _ = env.step(action.item())
        reward = torch.tensor([reward], device=device)

        # Observe new state
        last_screen = current_screen
        current_screen = get_screen()
        if not done:
            next_state = current_screen - last_screen
        else:
            next_state = None

        # Store the transition in memory
        memory.push(state, action, next_state, reward)

        # Move to the next state
        state = next_state

        # Perform one step of the optimization (on the target network)
        optimize_model()
        if done:
            episode_durations.append(t + 1)
            plot_durations()
            break
    # Update the target network
    if i_episode % TARGET_UPDATE == 0:
        target_net.load_state_dict(policy_net.state_dict())

print('Complete')
env.render()
env.close()
plt.ioff()
plt.show()