注意
轉至末尾 下載完整的示例程式碼。
迴圈 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 文件可以在這裡找到。
