注意
转至末尾 下载完整的示例代码。
循环 DQN：训练循环策略¶
如何在 TorchRL 的 actor 中整合 RNN
如何将此基于记忆的策略与回放缓冲区和损失模块一起使用
PyTorch v2.0.0
gym[mujoco]
tqdm
import tempfile
概述¶
基于记忆的策略不仅在观察结果部分可观察时至关重要，而且当需要考虑时间维度来做出明智的决策时也至关重要。
循环神经网络长期以来一直是基于记忆策略的流行工具。其思想是在两个连续步骤之间在内存中保留一个循环状态，并将其与当前观察结果一起作为策略的输入。
本教程演示了如何使用 TorchRL 将 RNN 整合到策略中。
主要学习内容
在 TorchRL 的 actor 中整合 RNN；
将此基于记忆的策略与回放缓冲区和损失模块一起使用。
在 TorchRL 中使用 RNN 的核心思想是使用 TensorDict 作为数据载体，在一步一步之间传递隐藏状态。我们将构建一个策略，该策略从当前 TensorDict 读取前一个循环状态，并将当前循环状态写入下一个状态的 TensorDict 中。
如图所示，我们的环境使用零初始化的循环状态填充 TensorDict，策略将其与观察结果一起读取以产生动作，同时还将循环状态用于下一步。当调用 step_mdp() 函数时，下一个状态的循环状态会被带到当前 TensorDict。让我们看看这在实践中是如何实现的。
如果您在 Google Colab 中运行此代码，请确保安装以下依赖项
!pip3 install torchrl
!pip3 install gym[mujoco]
!pip3 install tqdm
设置¶
import torch
import tqdm
from tensordict.nn import (
TensorDictModule as Mod,
TensorDictSequential,
TensorDictSequential as Seq,
)
from torch import nn
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
from torchrl.data.replay_buffers.samplers import SliceSampler
from torchrl.envs import (
Compose,
ExplorationType,
GrayScale,
InitTracker,
ObservationNorm,
Resize,
RewardScaling,
set_exploration_type,
StepCounter,
ToTensorImage,
TransformedEnv,
)
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import ConvNet, EGreedyModule, LSTMModule, MLP, QValueModule
from torchrl.objectives import DQNLoss, SoftUpdate
is_fork = multiprocessing.get_start_method() == "fork"
device = (
torch.device(0)
if torch.cuda.is_available() and not is_fork
else torch.device("cpu")
)
环境¶
像往常一样，第一步是构建我们的环境：它有助于我们定义问题并相应地构建策略网络。在本教程中，我们将运行一个基于像素的 CartPole gym 环境的单个实例，并进行一些自定义转换：转换为灰度、调整大小到 84x84、缩减奖励并归一化观察结果。
注意
StepCounter 转换是辅助性的。由于 CartPole 任务的目标是尽可能延长轨迹，因此计算步数有助于我们跟踪策略的性能。
为了本教程的目的，有两个转换很重要
InitTracker将通过在 TensorDict 中添加一个"is_init"布尔掩码来标记对reset()的调用，该掩码将跟踪哪些步骤需要重置 RNN 隐藏状态。TensorDictPrimer转换更技术性一些。它不是使用 RNN 策略所必需的。但是，它会指示环境（以及后续的收集器）预期存在一些额外的键。添加后，对 env.reset() 的调用将使用零初始化的张量填充初化器中指定的条目。知道策略需要这些张量，收集器将在收集过程中传递它们。最终，我们将把隐藏状态存储在回放缓冲区中，这将有助于我们在损失模块中启动 RNN 操作的计算（否则将使用 0 进行初化）。总之：不包含此转换不会对策略的训练产生巨大影响，但它会使循环键从收集到的数据和回放缓冲区中消失，从而导致训练效果略差。幸运的是，我们提供的LSTMModule包含一个帮助方法来创建这个转换，所以我们可以等到创建它时再处理！
env = TransformedEnv(
GymEnv("CartPole-v1", from_pixels=True, device=device),
Compose(
ToTensorImage(),
GrayScale(),
Resize(84, 84),
StepCounter(),
InitTracker(),
RewardScaling(loc=0.0, scale=0.1),
ObservationNorm(standard_normal=True, in_keys=["pixels"]),
),
)
一如既往，我们需要手动初化归一化常量
env.transform[-1].init_stats(1000, reduce_dim=[0, 1, 2], cat_dim=0, keep_dims=[0])
td = env.reset()
策略 (Policy)¶
我们的策略将有 3 个组成部分：一个 ConvNet 主干，一个 LSTMModule 内存层和一个浅层 MLP 模块，它将 LSTM 的输出映射到动作值。
卷积网络¶
我们构建了一个卷积网络，并用一个 torch.nn.AdaptiveAvgPool2d 进行了加装，该网络将输出压缩成一个大小为 64 的向量。 ConvNet 可以帮助我们实现这一点。
feature = Mod(
ConvNet(
num_cells=[32, 32, 64],
squeeze_output=True,
aggregator_class=nn.AdaptiveAvgPool2d,
aggregator_kwargs={"output_size": (1, 1)},
device=device,
),
in_keys=["pixels"],
out_keys=["embed"],
)
我们执行第一个模块来处理一批数据，以收集输出向量的大小。
n_cells = feature(env.reset())["embed"].shape[-1]
LSTM 模块¶
TorchRL 提供了一个专门的 LSTMModule 类，用于将 LSTMs 整合到您的代码库中。它是一个 TensorDictModuleBase 的子类：因此，它有一组 in_keys 和 out_keys，它们指示在模块执行期间应该预期读取和写入/更新的值。该类带有这些属性的可自定义预定义值，以方便其构建。
注意
使用限制：该类支持几乎所有的 LSTM 功能，如 dropout 或多层 LSTM。但是，为了遵守 TorchRL 的约定，此 LSTM 必须将 batch_first 属性设置为 True，这在 PyTorch 中**不是**默认值。但是，我们的 LSTMModule 改变了此默认行为，因此我们可以直接调用它。
此外，LSTM 不能将 bidirectional 属性设置为 True，因为这在在线环境中无法使用。在这种情况下，默认值是正确的。
lstm = LSTMModule(
input_size=n_cells,
hidden_size=128,
device=device,
in_key="embed",
out_key="embed",
)
让我们看看 LSTM Module 类，特别是它的 in 和 out_keys。
print("in_keys", lstm.in_keys)
print("out_keys", lstm.out_keys)
我们可以看到，这些值包含我们作为 in_key（和 out_key）指示的键以及循环键名称。out_keys 前面有一个“next”前缀，表明它们需要写入“next” TensorDict。我们使用此约定（可以通过传递 in_keys/out_keys 参数来覆盖）来确保调用 step_mdp() 会将循环状态移动到根 TensorDict，使其在下一次调用时可用于 RNN（参见介绍中的图）。
如前所述，我们还有一个可选的转换需要添加到我们的环境中，以确保循环状态被传递到缓冲区。 make_tensordict_primer() 方法正是为此目的而设计的。
env.append_transform(lstm.make_tensordict_primer())
就这样！我们可以打印环境以检查现在添加了初化器后一切看起来都很好。
print(env)
MLP¶
我们使用单层 MLP 来表示我们将用于策略的动作值。
mlp = MLP(
out_features=2,
num_cells=[
64,
],
device=device,
)
并用零填充偏置。
mlp[-1].bias.data.fill_(0.0)
mlp = Mod(mlp, in_keys=["embed"], out_keys=["action_value"])
使用 Q 值选择动作¶
我们策略的最后一部分是 Q 值模块。Q 值模块 QValueModule 将读取由我们的 MLP 生成的 "action_values" 键，并从中收集具有最大值的动作。我们唯一需要做的是指定动作空间，这可以通过传递字符串或动作规范来完成。这允许我们使用分类（有时称为“稀疏”）编码或其 one-hot 版本。
qval = QValueModule(action_space=None, spec=env.action_spec)
注意
TorchRL 还提供了一个包装器类 torchrl.modules.QValueActor，它将一个模块与 QValueModule 一起包装在 Sequential 中，就像我们在此处显式进行的那样。这样做的好处不大，而且过程不太透明，但最终结果将与我们在这里所做的类似。
我们现在可以将它们组合到一个 TensorDictSequential 中。
policy = Seq(feature, lstm, mlp, qval)
DQN 是一个确定性算法，探索是其中的关键部分。我们将使用一个 \(\epsilon\)-greedy 策略，其 epsilon 值为 0.2，并逐渐衰减到 0。通过调用 step()（见下面的训练循环）来实现此衰减。
exploration_module = EGreedyModule(
annealing_num_steps=1_000_000, spec=env.action_spec, eps_init=0.2
)
stoch_policy = TensorDictSequential(
policy,
exploration_module,
)
将模型用于损失¶
我们构建的模型已准备好在顺序设置中使用。但是，torch.nn.LSTM 类可以使用 cuDNN 优化的后端来更快地在 GPU 设备上运行 RNN 序列。我们不想错过这个加速训练循环的机会！
默认情况下，torchrl 损失在执行任何 LSTMModule 或 GRUModule 前向调用时都会使用此功能。如果您需要手动控制此功能，RNN 模块对上下文管理器/装饰器 set_recurrent_mode 敏感，该装饰器会处理底层 RNN 模块的行为。
由于我们还有一些未初始化的参数，因此在创建优化器等之前应该先初始化它们。
policy(env.reset())
DQN 损失¶
我们的 DQN 损失需要我们传递策略，并且再次传递动作空间。虽然这似乎是多余的，但很重要，因为我们想确保 DQNLoss 和 QValueModule 类是兼容的，但它们之间没有强依赖关系。
要使用 Double-DQN，我们需要 delay_value 参数，该参数将创建一个非可微分的网络参数副本，用作目标网络。
loss_fn = DQNLoss(policy, action_space=env.action_spec, delay_value=True)
由于我们使用的是 Double DQN，因此我们需要更新目标参数。我们将使用 SoftUpdate 实例来完成此工作。
updater = SoftUpdate(loss_fn, eps=0.95)
optim = torch.optim.Adam(policy.parameters(), lr=3e-4)
收集器和回放缓冲区¶
对于基于 RNN 的策略，我们需要采样连续转换序列，而不是独立转换。我们将使用 SliceSampler 来采样长度为 50 的轨迹切片。这确保了 LSTM 隐藏状态在训练期间能够通过序列正确传播。
该缓冲区将存储 1,000,000 个单独的转换，在采样时，我们将获得最多 50 个连续步骤的切片。在每个优化步骤（每次数据收集 16 次），我们将从轨迹切片中采样总计 200 个转换的批次。
我们将使用 LazyMemmapStorage 存储将数据保存在磁盘上，并将回放缓冲区直接传递给数据收集器，以便在收集数据时自动填充缓冲区。
注意
为了效率起见，我们这里只运行了数千次迭代。在实际环境中，总帧数应设置为 100 万。
buffer_scratch_dir = tempfile.TemporaryDirectory().name
rb = TensorDictReplayBuffer(
storage=LazyMemmapStorage(1_000_000, scratch_dir=buffer_scratch_dir),
sampler=SliceSampler(
slice_len=50, end_key=("next", "done"), cache_values=True, strict_length=False
),
batch_size=200,
prefetch=10,
transform=lambda td: td.to(device),
)
collector = SyncDataCollector(
env,
stoch_policy,
frames_per_batch=50,
total_frames=200,
storing_device="cpu",
replay_buffer=rb,
)
训练循环¶
由于我们将回放缓冲区传递给了收集器，因此收集器不再直接生成数据 - 而是自动填充回放缓冲区。我们通过收集器进行迭代以触发数据收集，然后从缓冲区采样进行训练。为了跟踪进度，我们将每 50 次数据收集迭代运行一次策略，并在训练后绘制结果。
utd = 16
pbar = tqdm.tqdm(total=collector.total_frames)
longest = 0
traj_lens = []
for i, _ in enumerate(collector):
pbar.update(collector.frames_per_batch)
# Only start training once we have enough data in the buffer
if len(rb) < 1000:
continue
for j in range(utd):
s = rb.sample()
if i == 0 and j == 0:
# Let's print the first sample to see the data structure
print(
"Let us print the first batch of sampled data from the replay buffer.\n"
"Pay attention to the key names which will reflect what can be found in this data structure, "
"in particular: the output of the QValueModule (action_values, action and chosen_action_value),"
"the 'is_init' key that will tell us if a step is initial or not, and the "
"recurrent_state keys.\n",
s,
)
loss_vals = loss_fn(s)
loss_vals["loss"].backward()
optim.step()
optim.zero_grad()
pbar.set_description(
f"buffer_size: {len(rb)}, loss_val: {loss_vals['loss'].item(): 4.4f}"
)
exploration_module.step(collector.frames_per_batch)
updater.step()
with set_exploration_type(ExplorationType.DETERMINISTIC), torch.no_grad():
rollout = env.rollout(10000, stoch_policy)
traj_lens.append(rollout.get(("next", "step_count")).max().item())
让我们绘制我们的结果。
if traj_lens:
from matplotlib import pyplot as plt
plt.plot(traj_lens)
plt.xlabel("Test collection")
plt.title("Test trajectory lengths")
结论¶
我们已经了解了如何在 TorchRL 中将 RNN 整合到策略中。您现在应该能够
创建一个充当
TensorDictModule的 LSTM 模块。通过
InitTracker转换指示 LSTM 模块需要重置。将此模块整合到策略和损失模块中。
确保收集器了解循环状态条目，以便它们可以与其余数据一起存储在回放缓冲区中。
延伸阅读¶
TorchRL 文档可以在这里找到。
