REINFORCEMENT LEARNING (DQN) TUTORIAL

作者: Adam Paszke ,Mark Towers

本教程展示了如何使用 PyTorch 在 Gymnasium 的 CartPole-v1 任务上训练深度 Q 学习 （DQN） 代理。

cartpole

当代理观察环境的当前状态并选择操作时，环境将转换为新状态，并且还会返回指示操作后果的奖励。在此任务中，每个增量时间步的奖励为 +1，如果杆子掉落太远或推车偏离中心超过 2.4 个单位，环境就会终止。这意味着表现更好的场景将运行更长时间，从而积累更大的回报。

CartPole 任务的设计使代理的输入是 4 个表示环境状态（位置、速度等）的真实值。我们获取这 4 个没有任何缩放的输入，并将它们传递到一个具有 2 个输出的小型全连接网络，每个操作一个。网络经过训练，以预测给定输入状态的每个操作的预期值。然后选择具有最高期望值的操作。

任务

代理必须在两个动作之间做出决定 - 向左或向右移动推车 - 以便连接到它的杆子保持直立。您可以在 Gymnasium的网站上找到有关环境和其他更具挑战性环境的更多信息。

包

首先，让我们导入所需的包。首先，我们需要环境体育馆，使用pip安装。这是原始OpenAI Gym项目的分支，自Gym v0.19以来由同一团队维护。

pip install gymnasium[classic_control]

我们将使用PyTorch的如下内容:

• 神经网络 (torch.nn)

• 优化 (torch.optim)

• 自动微分 (torch.autograd)

import gymnasium as gym
import math
import random
import matplotlib
import matplotlib.pyplot as plt
from collections import namedtuple, deque
from itertools import countimport torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as Fenv = gym.make("CartPole-v1")# set up matplotlib
is_ipython = 'inline' in matplotlib.get_backend()
if is_ipython:from IPython import displayplt.ion()# if gpu is to be used
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

记忆回放

我们将使用经验记忆回放来训练我们的 DQN。它存储代理观察到的转换，允许我们稍后重用这些数据。通过从中随机采样，构建批次的转换是去相关的。已经表明，这极大地稳定和改进了DQN训练过程。

为此，我们需要两个类：

Transition ——即转移，一个命名元组，表示我们环境中的单个过渡。它本质上将 (state,action) 即（状态，操作）对映射到它们的 (next_state, reward) 即（下一状态，奖励）结果，状态是接下来描述的屏幕差异图像。

ReplayMemory ——即记忆回放—— 有界大小的循环缓冲区，用于保存最近观察到的转换。它还实现了一个 .sample（） 方法，用于选择随机批次的转换进行训练。

Transition = namedtuple('Transition',('state', 'action', 'next_state', 'reward'))class ReplayMemory(object):def __init__(self, capacity):self.memory = deque([], maxlen=capacity)def push(self, *args):"""Save a transition"""self.memory.append(Transition(*args))def sample(self, batch_size):return random.sample(self.memory, batch_size)def __len__(self):return len(self.memory)

DQN算法

我们的环境是确定性的，因此为了简单起见，这里介绍的所有方程也是确定性的。在强化学习文献中，它们还将包含对环境中随机转换的期望。

我们的目标是训练一个策略，试图最大化折扣的累积奖励:

其中Rt0也称为回报。折扣γ应为 0 到 1 之间的常量，以确保总和收敛。较低的γ使得来自不确定的遥远未来的奖励对我们的代理来说不如它可以有相当信心的在不久的将来的奖励重要。它还鼓励代理收集在时间上相对遥远的时间更近的时间上的等价的奖励。

Q-learning背后的主要思想是，如果我们有一个函数Q*:State×Action→R，它可以告诉我们我们的回报是多少，如果我们要在给定状态下采取行动，那么我们可以很容易地构建一个最大化奖励的策略：

但是，我们对世界一无所知，因此无法访问Q*。但是，由于神经网络是通用函数逼近器，我们可以简单地创建一个并训练它类似于 Q *。

对于我们的训练更新规则，我们将使用一个事实，即某些策略的每个 Q 函数都遵循贝尔曼方程：

等式两边的差称为时差误差，δ：

为了尽量减少此错误，我们将使用 Huber 损失。当误差较小时，Huber 损失的作用类似于均方误差，但当误差较大时，Huber 损失的作用类似于平均绝对误差 - 这使得当 Q 的估计值非常嘈杂时，它对异常值更稳健。我们通过从回放记忆中采样的一批转换 B 来计算这一点：

Q网络

我们的模型将是一个卷积神经网络（译者注：实现代码实际是全连接网络），它考虑了当前和先前屏幕小快之间的差异。它有两个输出，分别表示Q(s,left)和Q(s,right)（其中s是网络的输入）。实际上，网络试图预测在给定当前输入的情况下，采取每项行动的预期回报

class DQN(nn.Module):def __init__(self, n_observations, n_actions):super(DQN, self).__init__()self.layer1 = nn.Linear(n_observations, 128)self.layer2 = nn.Linear(128, 128)self.layer3 = nn.Linear(128, n_actions)# Called with either one element to determine next action, or a batch# during optimization. Returns tensor([[left0exp,right0exp]...]).def forward(self, x):x = F.relu(self.layer1(x))x = F.relu(self.layer2(x))return self.layer3(x)

训练

超参数和实用程序

此单元实例化我们的模型及其优化器，并定义一些实用程序：

• select_action - 将根据 epsilon 贪婪策略选择操作。简单地说，我们有时会使用我们的模型来选择动作，有时我们只会统一采样一个。选择随机操作的概率将从EPS_START开始，并将呈指数衰减到EPS_END。EPS_DECAY控制衰减的速度。

• plot_durations - 绘制剧集持续时间以及过去 100 集的平均值（官方评估中使用的度量）的助手。该图将位于包含主训练循环的单元格下方，并将在每集后更新。

# BATCH_SIZE is the number of transitions sampled from the replay buffer
# GAMMA is the discount factor as mentioned in the previous section
# EPS_START is the starting value of epsilon
# EPS_END is the final value of epsilon
# EPS_DECAY controls the rate of exponential decay of epsilon, higher means a slower decay
# TAU is the update rate of the target network
# LR is the learning rate of the AdamW optimizer
BATCH_SIZE = 128
GAMMA = 0.99
EPS_START = 0.9
EPS_END = 0.05
EPS_DECAY = 1000
TAU = 0.005
LR = 1e-4# Get number of actions from gym action space
n_actions = env.action_space.n
# Get the number of state observations
state, info = env.reset()
n_observations = len(state)policy_net = DQN(n_observations, n_actions).to(device)
target_net = DQN(n_observations, n_actions).to(device)
target_net.load_state_dict(policy_net.state_dict())optimizer = optim.AdamW(policy_net.parameters(), lr=LR, amsgrad=True)
memory = ReplayMemory(10000)steps_done = 0def select_action(state):global steps_donesample = random.random()eps_threshold = EPS_END + (EPS_START - EPS_END) * \math.exp(-1. * steps_done / EPS_DECAY)steps_done += 1if sample > eps_threshold:with torch.no_grad():# t.max(1) will return the largest column value of each row.# second column on max result is index of where max element was# found, so we pick action with the larger expected reward.return policy_net(state).max(1)[1].view(1, 1)else:return torch.tensor([[env.action_space.sample()]], device=device, dtype=torch.long)episode_durations = []def plot_durations(show_result=False):plt.figure(1)durations_t = torch.tensor(episode_durations, dtype=torch.float)if show_result:plt.title('Result')else:plt.clf()plt.title('Training...')plt.xlabel('Episode')plt.ylabel('Duration')plt.plot(durations_t.numpy())# Take 100 episode averages and plot them tooif 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 updatedif is_ipython:if not show_result:display.display(plt.gcf())display.clear_output(wait=True)else:display.display(plt.gcf())

训练循环

最后，用于训练模型的代码。

在这里，您可以找到一个执行优化步骤的optimize_model函数。它首先对一批进行采样，将所有张量连接成一个张量，计算Q(st,at) 和 V(st+1)=maxa(st+1,a)，并将它们组合成我们的损失。根据定义，如果 s 是终端状态，我们设置 V(s)=0。我们还使用目标网络来计算 V(st+1) 以增加稳定性。目标网络在每一步都使用由之前定义的超参数 TAU 控制的软更新进行更新。

def optimize_model():if len(memory) < BATCH_SIZE:returntransitions = memory.sample(BATCH_SIZE)# Transpose the batch (see https://stackoverflow.com/a/19343/3343043 for# detailed explanation). This converts batch-array of Transitions# to Transition of batch-arrays.batch = Transition(*zip(*transitions))# Compute a mask of non-final states and concatenate the batch elements# (a final state would've been the one after which simulation ended)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_stateif 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. These are the actions which would've been taken# for each batch state according to policy_netstate_action_values = policy_net(state_batch).gather(1, action_batch)# Compute V(s_{t+1}) for all next states.# Expected values of actions for non_final_next_states are computed based# on the "older" target_net; selecting their best reward with max(1)[0].# This is merged based on the mask, such that we'll have either the expected# state value or 0 in case the state was final.next_state_values = torch.zeros(BATCH_SIZE, device=device)with torch.no_grad():next_state_values[non_final_mask] = target_net(non_final_next_states).max(1)[0]# Compute the expected Q valuesexpected_state_action_values = (next_state_values * GAMMA) + reward_batch# Compute Huber losscriterion = nn.SmoothL1Loss()loss = criterion(state_action_values, expected_state_action_values.unsqueeze(1))# Optimize the modeloptimizer.zero_grad()loss.backward()# In-place gradient clippingtorch.nn.utils.clip_grad_value_(policy_net.parameters(), 100)optimizer.step()

下面，您可以找到主要的训练循环。一开始我们重置环境并获得初始状态张量。然后，我们对一个动作进行采样，执行它，观察下一个状态和奖励（始终为 1），并优化我们的模型一次。当训练的一集结束时（我们的模型失败），我们重新启动循环。

下面，如果 GPU 可用，num_episodes设置为 600，否则安排 50 集，因此训练不会花费太长时间。然而，50集不足以在 cartpole上观察到良好的表现。您应该看到模型在 600 个训练集中不断达到 500 步。训练 RL 代理可能是一个嘈杂的过程，因此如果不观察到收敛，重新启动训练可以产生更好的结果。

if torch.cuda.is_available():num_episodes = 600
else:num_episodes = 50for i_episode in range(num_episodes):# Initialize the environment and get it's statestate, info = env.reset()state = torch.tensor(state, dtype=torch.float32, device=device).unsqueeze(0)for t in count():action = select_action(state)observation, reward, terminated, truncated, _ = env.step(action.item())reward = torch.tensor([reward], device=device)done = terminated or truncatedif terminated:next_state = Noneelse:next_state = torch.tensor(observation, dtype=torch.float32, device=device).unsqueeze(0)# Store the transition in memorymemory.push(state, action, next_state, reward)# Move to the next statestate = next_state# Perform one step of the optimization (on the policy network)optimize_model()# Soft update of the target network's weights# θ′ ← τ θ + (1 −τ )θ′target_net_state_dict = target_net.state_dict()policy_net_state_dict = policy_net.state_dict()for key in policy_net_state_dict:target_net_state_dict[key] = policy_net_state_dict[key]*TAU + target_net_state_dict[key]*(1-TAU)target_net.load_state_dict(target_net_state_dict)if done:episode_durations.append(t + 1)plot_durations()breakprint('Complete')
plot_durations(show_result=True)
plt.ioff()
plt.show()

操作是随机选择的，也可以根据策略选择，从gymnasium环境中获取下一步样本。我们将结果记录在回放记忆中，并在每次迭代时运行优化步骤。优化从回放记忆中随机选取一个批次来训练新策略。“旧”target_net也用于优化以计算预期的 Q 值。每一步都对其权重进行软更新。

原文

• Reinforcement Learning (DQN) Tutorial — PyTorch Tutorials 1.13.1+cu117 documentation