pytorch实现CartPole-v1任务的DQN代码

DQN

使用PyTorch在OpenAI Gym上的CartPole-v1任务上训练深度Q学习（DQN）智能体

任务

CartPole-v1环境中，手推车上面有一个杆，手推车沿着无摩擦的轨道移动。 通过对推车施加+1或-1的力来控制系统。 钟摆最开始为直立状态，训练的目的是防止其跌落。 杆保持直立的每个时间步长都提供+1的奖励。 当杆与垂直线的夹角超过15度时，或者推车从中心移出2.4个单位以上时，训练结束。智能体必须在两个动作之间做出决定-向左或向右移动推车-以便使与之相连的杆子保持直立。表现更好的方案将持续更长的时间，从而积累更大的回报

CartPole任务中智能体的输入是代表环境状态（位置，速度等）的4个实际值。神经网络可以通过查看场景来解决任务，因此以推车为中心的一部分屏幕作为输入。

包

首先导入需要的包，先安装针对环境的gym，还需要来自pytorch的以下模块

• 神经网络 (torch.nn)
• 优化 (torch.optim)
• 自动微分 (torch.autograd)
• 视觉任务的工具 (torchvision)

import gym
import math
import random
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple, deque
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-v1').unwrapped

# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:
    from IPython import display
# matplotlib.get_backend()=inline即当前后端为ipykernel.pylab.backend_inline
# is_ipython=True

plt.ion()
# plt.ion()函数能使matplotlib的显示模式转换为交互（interactive）模式。

# if gpu is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

Matplotlib在设计上对用户编写的绘图代码和对不同输出形式的处理方法进行了隔离，出现了前端（frontend）和后端的概念（backend）。后端可以认为就是不同输出形式的处理功能，前端可以认为就是用户所要绘制的图像。就像Web开发中的前后端分离一样，用户只用关心如何绘图即可，Matplotlib会根据用户选择的后端进行输出。这样相同的前端绘图代码，就可以便捷地实现各种绘图输出。
matplotlib.get_backend()可以得到当前使用的后端名称

1.Replay Memory

DQN经验回放减少了数据间的关联性，智能体观察到的转换存储到经验池中，之后从中随机采样进行训练

为此需要两个类:

• Transition - 表示环境中单个过渡的命名元组，即将（state，action）对映射到其（next_state, reward）结果。此处状态是屏幕差异图像
• ReplayMemory - 有容量限制的经验回放池，用于保存最近观察到的转换。它还实现了.sample()方法，用于选择随机的过渡批量进行训练。

# 定义Transition是一个包含四个属性的namedtuple类型
Transition = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward'))

# 定义经验回放类
class ReplayMemory(object):

    def __init__(self, capacity):
        self.memory = deque([],maxlen=capacity) # 经验池memory是个最大容量为capacity的队列

    def push(self, *args):
        """Save a transition"""
        self.memory.append(Transition(*args)) # 存储新的Transition

    def sample(self, batch_size):
        return random.sample(self.memory, batch_size) # 从memory中随机选择batch_size个元素返回

    def __len__(self):
        return len(self.memory) # 返回当前memory的数量

namedtuples是继承自tuple的子类，namedtuple创建一个和tuple类似的对象，而且对象拥有可访问的属性。

# 示例
from collections import namedtuple
# 定义一个namedtuple类型User，并包含name，sex和age属性。
User = namedtuple('User',['name','sex','age'])
# 创建一个User对象
user = User._make(['lenmo','male','20'])
# 输出user的姓名
print(user.name)

lenmo

python3中的列表可以当成堆栈和队列使用，deque创建一个队列

# 示例
from collections import deque
queue = deque(['a','b','c'])
queue.append('d')
queue.popleft()
queue

deque(['b', 'c', 'd'])

*args：当传入的参数个数未知，且不需要知道参数名称时

**args：当传入的参数个数未知，但需要知道参数的名称时(类似字典）

2.CNN网络

模型是一个卷积神经网络，输入当前屏幕补丁和之前屏幕补丁的差异。它有两个输出，分别代表$Q(s,\mathrm{l}\mathrm{e}\mathrm{f}\mathrm{t})$和$Q(s,\mathrm{r}\mathrm{i}\mathrm{g}\mathrm{h}\mathrm{t})$，其中$s$是网络的输入，网络在尝试预测在当前输入下执行每个操作的预期收益

# outputs=动作数量，CNN输入每个状态特征值，输出每个state-action的值函数
class DQN(nn.Module):
    
    def __init__(self,h,w,outputs): 
        super(DQN,self).__init__()
        self.conv1 = nn.Conv2d(3,16,kernel_size=5,stride=2)
        # 输入(b,3,h,w)，16个过滤器（3，5，5）卷积计算，输出（b,16,h1,w1)
        self.bn1 = nn.BatchNorm2d(16)
        # Batch 归一化
        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)
        
        # 计算卷积后的尺寸用于后续平铺到全连接层
        def conv2d_size_out(size,kernel_size=5,stride=2):
            return (size-kernel_size)//stride + 1
        convw = conv2d_size_out(conv2d_size_out(conv2d_size_out(w)))
        convh = conv2d_size_out(conv2d_size_out(conv2d_size_out(h)))
        
        # 最后一个卷积层的输出平铺作为全连接层的输入
        linear_input_size = convh * convw * 32
        self.head = nn.Linear(linear_input_size,outputs)

    def forward(self,x):
        x = x.to(device)
        x = F.relu(self.bn1(self.conv1(x))) # 卷积->batch归一化->激活
        x = F.relu(self.bn2(self.conv2(x)))
        x = F.relu(self.bn3(self.conv3(x)))
        # x.view相当于reshape，x.size(0)是batch维度,x.view(a,-1)表示a行，列根据原先尺寸决定
        return self.head(x.view(x.size(0),-1)) 

torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None)：在4D输入上应用 Batch归一化(带有通道尺寸的 mini-batch 2D输入)

num_features：输入（N,C,H,W)中的C

# 示例
M = nn.BatchNorm2d(3)
input = torch.randn(2,3,35,45)
output = M(input)
print(input.size(0))
input.view(input.size(0),-1)

2

tensor([[ 1.8808,  0.3987,  1.1235,  ...,  0.0297,  0.0279, -1.2829],
        [ 2.0374, -0.9642, -2.1847,  ...,  0.2094,  0.4277,  0.4624]])

3.输入提取

以下代码是用于从环境中提取和处理渲染图像的工具，使用了torchvision包，可以轻松组成图像转换，运行单位就会显示提取的实例补丁

图像处理操作

# ToPILImage将张量或者多维数组变成PIL图像
# Resize将输入图像尺寸变成给定尺寸，
# ToTensor将PIL图像或者多维数组变成张量
resize = T.Compose([T.ToPILImage(),
                    T.Resize(40,interpolation=Image.CUBIC),
                    T.ToTensor()])

获取小车的屏幕坐标

# 小车世界是（-4.8，4.8），小车有效世界是（-2.4，2.4）即小车移动超过2.4会终止episode，env.x_threshold=2.4
# 屏幕世界是高400*宽600，屏幕世界中心（0，0）在左下角
# cartenv源码设置self.state = (x, x_dot, theta, theta_dot)，env.state[0]是cart位置
# 世界坐标（0，0）对应屏幕坐标中间（300，0）位置，因此需要加上屏幕宽度一半

def get_cart_location(screen_width):
    world_width = env.x_threshold * 2 # world_width是小车世界可移动宽度
    scale = screen_width / world_width # scale是屏幕世界与小车世界宽度的比值（转屏幕系数）
    return int(env.state[0] * scale + screen_width / 2.0) # 计算小车的屏幕世界坐标

获取处理后的屏幕图像张量

def get_screen():
    # gym中的屏幕是400*600*3，有时会更大比如800*1200*3，先将其变成torch里的顺序（C,H,W）即3*400*600
    screen = env.render(mode='rgb_array').transpose(2,0,1)
    # 获取屏幕世界宽度600和高度400
    _,screen_height,screen_width = screen.shape
    # 小车初始状态位置在屏幕高40%（400X0.4=160）到80%（400X0.8=320）之间，高度按照160-320截取
    screen = screen[:,int(screen_height*0.4):int(screen_height*0.8)]
    # view_width 60%屏幕宽度
    view_width = int(screen_width * 0.6)
    cart_location = get_cart_location(screen_width)
    
    # 当前小车位置偏左即左边没有30%的空间，截取最左侧60%
    if cart_location < view_width // 2:
        slice_range = slice(view_width)
    # 当前小车位置偏右即右边没有30%的空间，截取最右侧60%
    elif cart_location > (screen_width-view_width // 2):
        slice_range = slice(-view_width,None)
    # 当前小车位置左右均有30%空间，保留小车左右各30%的位置
    else:
        slice_range = slice(cart_location - view_width // 2,cart_location + view_width // 2)
    screen = screen[:,:,slice_range]
    
    # python数组默认行优先存储（行连续），列不连续，经过行的slice操作后，会使数组行没有连续性
    # ascontiguousarray函数将一个内存不连续存储的数组转换为内存连续存储的数组，使得运行速度更快。
    screen = np.ascontiguousarray(screen,dtype=np.float32) / 225
    screen = torch.from_numpy(screen) # 从numpy.ndarray创建一个张量
    return resize(screen).unsqueeze(0) # 使用unsqueeze(0)扩展0维即batch维度变成（BCHW)

显示一张屏幕图像样例

env.reset() # 重置环境
plt.figure()
# 从get_screen中返回的是4维张量，使用squeeze去掉batch维度
# torch.Tensor.permute():将tensor维度换位，变成(HWC)
plt.imshow(get_screen().cpu().squeeze(0).permute(1,2,0).numpy(),interpolation='none')
plt.title("example extracted screen")
plt.show()

4.训练

超参数和工具

实例化模型和优化器，定义工具

• 选择动作:根据 ε 贪婪策略选择一个动作,选择随机动作的可能性将从EPS_START开始，并朝EPS_END呈指数衰减。 EPS_DECAY控制衰减率。
• 动态显示:绘制episodes持续时间和最近100个episodes的均值，每个episodes之后更新

定义TD目标网络和动作值函数网络

# 经过get_screen()的裁剪和处理后，返回的初始屏幕即图像是3*40*90
init_screen = get_screen()
_,_,screen_height,screen_width = init_screen.shape

# gym动作空间数量=2（left,right）
n_actions = env.action_space.n

# 定义策略网络和目标网络
policy_net = DQN(screen_height,screen_width,n_actions).to(device)
target_net = DQN(screen_height,screen_width,n_actions).to(device)
# load_state_dict()函数用于将预训练的参数权重加载到新的模型之中
# state_dict()将每一层与它的对应参数建立映射关系，python字典存储参数
# TD目标网络的参数按步长根据策略网络的参数更新
target_net.load_state_dict(policy_net.state_dict())
# TD目标网络为eval()模式，即告诉pytorch该网络不参与训练
target_net.eval()

DQN(
  (conv1): Conv2d(3, 16, kernel_size=(5, 5), stride=(2, 2))
  (bn1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(2, 2))
  (bn2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (conv3): Conv2d(32, 32, kernel_size=(5, 5), stride=(2, 2))
  (bn3): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (head): Linear(in_features=512, out_features=2, bias=True)
)

探索和利用（选择动作）ε-decreasing strategy

optimizer = optim.RMSprop(policy_net.parameters()) # 使用RMSprop优化方法
memory = ReplayMemory(1000) # 设置容量为1000的经验池

step_done = 0

# 参数设置
BATCH_SIZE = 128
GAMMA = 0.999
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 200
TARGET_UPDATE = 10

def select_action(state):
    global step_done
    sample = random.random() # 随机生成(0,1)值与eps比较
    eps_threshold = EPS_END + (EPS_START-EPS_END) * math.exp(-1. * step_done / EPS_DECAY) # eps随着训练步长下降
    step_done += 1
    if sample > eps_threshold:
        with torch.no_grad():
            # t.max(1)返回两列：每行最大的值+最大值的列索引
            # 返回Q值最大的动作索引0/1
            # 此处只有两个动作，因此policy_net(state)输出张量尺寸(1,2)
            return policy_net(state).max(1)[1].view(1,1) 
    else:
        # 随机选择动作，生成0/1
        return torch.tensor([[random.randrange(n_actions)]],device=device,dtype=torch.long)

ret = x.unfold(dim, size, step)

• dim：int，表示需要展开的维度(可以理解为窗口的方向)
• size：int，表示滑动窗口大小
• step：int，表示滑动窗口的步长

# 示例
import torch
x=torch.randn(4,3)
print(x)
x=x.unfold(0,2,1)
x

tensor([[-0.6728,  1.4748, -1.2390],
        [ 0.3977, -1.2479,  0.2498],
        [-1.1414, -1.8391,  1.3870],
        [ 0.1753,  0.5563, -0.2556]])

tensor([[[-0.6728,  0.3977],
         [ 1.4748, -1.2479],
         [-1.2390,  0.2498]],

        [[ 0.3977, -1.1414],
         [-1.2479, -1.8391],
         [ 0.2498,  1.3870]],

        [[-1.1414,  0.1753],
         [-1.8391,  0.5563],
         [ 1.3870, -0.2556]]])

绘制图像

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())
    # 取100个episode均值并绘图
    if len(durations_t) >= 100:
        means = durations_t.unfold(0,100,1).mean(1).view(-1)
        # torch.cat将两个张量（tensor）拼接在一起
        means = torch.cat((torch.zeros(99),means))
        plt.plot(means.numpy())
        
    plt.pause(0.001) # 暂停0.001s使点更新
    if is_ipython:
        display.clear_output(wait=True)
        display.display(plt.gcf())

训练循环（优化模型+主循环）

以下为训练模型的代码，即执行优化步骤的optimize_model函数，首先采样一个batch，将所有张量连接成一个张量，计算$Q(s_t,a_t)$和$V(s_{t+1})={\max }_{a}Q(s_{t+1},a)$，将其合并为损失。根据定义，如果$s$为终端状态，则设置$V(s)=0$。 使用目标网络来计算$V(s[t+1])$

def optimize_model():
    if len(memory) < BATCH_SIZE:
        return
    transitions = memory.sample(BATCH_SIZE) # 从经验池中随机采样BATCH_SIZE
    # 将batch_size个四元组，转换成，四个元祖，每个元祖有batch_size个项
    # 如选择2个样本(1,1,1,1)和(2,2,2,2)，转换后batch=Transition(state=(1,2),action=(1,2),next_state=(1,2),reward=(1,2))
    batch = Transition(*zip(*transitions))
    
    # lambda s:s is not None:输入s,输出一个bool值判断s是否为空
    # (map(lambda s:s is not None,batch.next_state)为每个样本的next_state执行lambda操作即s=batch.next_state
    # 整句代码是输出每个样本batch.next_state的True/False，区分终止状态和非终止状态
    non_final_mask = torch.tensor(tuple(map(lambda s:s is not None,batch.next_state)),device=device,dtype=torch.bool)
    # non_final_next_state存储非终止状态值
    non_final_next_state = 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)
    
    # gather(dim,index):按照index取值，dim决定索引维度
    state_action_values = policy_net(state_batch).gather(1,action_batch)
    
    next_state_values = torch.zeros(BATCH_SIZE,device=device)
    # 更新下一非空状态的Q值，选择值函数最大的动作对应的state-action value
    next_state_values[non_final_mask] = target_net(non_final_next_state).max(1)[0].detach()
    # 计算预期Q值（贝尔曼方程）
    expeacted_state_action_values = (next_state_values * GAMMA) + reward_batch
    
    # 计算TD误差
    criterion = nn.SmoothL1Loss()
    loss = criterion(state_action_values,expeacted_state_action_values.unsqueeze(1))
    
    # 优化模型
    optimizer.zero_grad()
    loss.backward()
    for param in policy_net.parameters():
        param.grad.data.clamp_(-1,1) # 梯度截断，防止出现梯度爆炸，将梯度约束在(-1,1)间
    optimizer.step()

# 示例
from collections import namedtuple
Transition11 = namedtuple('Transition',
                        ('state', 'action', 'next_state', 'reward'))
tran1 = Transition11._make([1,2,3,4])
tran2 = Transition11._make([11,22,33,44])
tran3 = Transition11._make([11,22,0,44])
batch1 = Transition11(*zip(tran1,tran2,tran3))
x = torch.tensor(tuple(map(lambda s:s is not None,batch1.next_state)),dtype=torch.bool)
y = torch.tensor([s for s in batch1.next_state if s is not None])
z = torch.zeros(3)
bs = torch.tensor(batch1.state)
print(x,y,z)
print(z[x])

tensor([True, True, True]) tensor([ 3, 33,  0]) tensor([0., 0., 0.])
tensor([0., 0., 0.])

以下是主循环代码，开始阶段重置环境，初始化state张量，然后随机选择一个动作并执行，观察下一个屏幕和奖励（总是1），并且优化一次模型。当episode结束时或者模型失败，重置循环。此处的num_episodes设置较小，一般超过300才会出现有意义的改进

num_episodes = 50
for i_episode in range(num_episodes):
    env.reset()
    last_screen = get_screen()
    current_screen = get_screen()
    state = current_screen - last_screen
    for t in count():
        # 选择动作并执行
        action = select_action(state)
        # 执行动作，env.step返回np.array(self.state, dtype=np.float32), reward, done, {}
        # done指小车距离或者角度超出规定范围，训练结束
        _,reward,done,_ = env.step(action.item())
        # env返回的reward是一个变量值，将其变成张量
        reward = torch.tensor([reward],device=device)

        # 观察下一个状态
        last_screen = current_screen
        current_screen = get_screen()
        if not done:
            next_state = current_screen - last_screen
        else:
            next_state = None

        # 将transition存储在经验池中
        memory.push(state,action,next_state,reward)

        # 转移到下一个状态
        state = next_state


        # 策略网络的一步优化
        optimize_model()
        if done:
            episode_durations.append(t + 1)
            plot_durations()
            break
    if i_episode % TARGET_UPDATE == 0: # TARGET_UPDATE个episode后更新目标网络的参数
        target_net.load_state_dict(policy_net.state_dict())

print('complete')
env.render() # 图像引擎
env.close() # 关闭环境
plt.ioff() # 显示图像前关掉交互模式
plt.show()

<Figure size 432x288 with 0 Axes>
complete
<Figure size 432x288 with 0 Axes>

说明总体结果数据流的图