迴圈 DQN：訓練迴圈策略

建立時間：2023 年 11 月 8 日 | 最後更新：2025 年 1 月 27 日 | 最後驗證：未驗證

作者：Vincent Moens

你將學到什麼

• 如何在 TorchRL 中將 RNN 整合到 actor 中

• 如何將基於記憶體的策略與回放緩衝區和損失模組一起使用

先決條件

• PyTorch v2.0.0

• gym[mujoco]

• tqdm

概述

基於記憶體的策略不僅在觀測部分可見時至關重要，而且在必須考慮時間維度以做出明智決策時也至關重要。

迴圈神經網路長期以來一直是基於記憶體策略的熱門工具。其思想是在兩個連續步驟之間將迴圈狀態儲存在記憶體中，並將其與當前觀測一起用作策略的輸入。

本教程展示瞭如何在 TorchRL 中使用 RNN 來實現策略。

關鍵學習要點

• 在 TorchRL 中將 RNN 整合到 actor 中；

• 將基於記憶體的策略與回放緩衝區和損失模組一起使用。

在 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 as Seq
from torch import nn
from torchrl.collectors import SyncDataCollector
from torchrl.data import LazyMemmapStorage, TensorDictReplayBuffer
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 環境的單個基於畫素的例項，並帶有一些自定義 Transforms：轉換為灰度圖、調整大小到 84x84、縮小獎勵並標準化觀測值。

注意

StepCounter transform 是輔助性的。由於 CartPole 任務的目標是使軌跡儘可能長，計算步數可以幫助我們跟蹤策略的效能。

對於本教程的目的，有兩個重要的 Transforms

• InitTracker 將透過在 TensorDict 中新增一個 "is_init" 布林掩碼來標記對 reset() 的呼叫，該掩碼將跟蹤哪些步驟需要重置 RNN 隱藏狀態。

• TensorDictPrimer transform 技術性稍強。它不是使用 RNN 策略所必需的。然而，它指示環境（以及隨後的 collector）預期一些額外的鍵。新增後，呼叫 env.reset() 將使用歸零張量填充 Primer 中指示的條目。知道這些張量是策略所期望的，collector 在收集過程中會傳遞它們。最終，我們會將隱藏狀態儲存在回放緩衝區中，這將有助於我們在損失模組中引導 RNN 操作的計算（否則將以 0 初始化）。總之：不包含此 transform 不會對策略的訓練產生巨大影響，但會導致迴圈鍵從收集到的資料和回放緩衝區中消失，這反過來會導致訓練稍微欠佳。幸運的是，我們提出的 LSTMModule 配備了一個幫助方法來構建正是這個 transform，所以我們可以等到構建它之後再來！

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()

策略

我們的策略將有 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 Module

TorchRL 提供了一個專門的 LSTMModule 類，用於將 LSTM 整合到你的程式碼庫中。它是 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” 字首，表明它們需要寫入“下一個” TensorDict 中。我們使用此約定（可以透過傳遞 in_keys/out_keys 引數來覆蓋）來確保呼叫 step_mdp() 會將迴圈狀態移動到根 TensorDict，使其在下一次呼叫時可用於 RNN（參見簡介中的圖）。

如前所述，我們還有一個可選的 transform 要新增到我們的環境中，以確保迴圈狀態傳遞到緩衝區。make_tensordict_primer() 方法正是做了這件事

env.append_transform(lstm.make_tensordict_primer())

就這樣！我們現在可以列印環境，檢查新增 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" 鍵，並從中獲取具有最大值的動作。我們唯一需要做的是指定動作空間，可以透過傳遞字串或 action-spec 來完成。這使我們可以使用分類（有時稱為“稀疏”）編碼或其獨熱版本。

qval = QValueModule(spec=env.action_spec)

注意

TorchRL 還提供了一個包裝類 torchrl.modules.QValueActor，它將模組與 QValueModule 一起包裝到 Sequential 中，就像我們在這裡明確做的那樣。這樣做的好處不大，而且過程不太透明，但最終結果與我們在此處所做類似。

我們現在可以將這些元件組合到一個 TensorDictSequential 中

stoch_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 = Seq(
    stoch_policy,
    exploration_module,
)

使用模型計算損失

我們構建的模型非常適合在順序設定中使用。但是，類 torch.nn.LSTM 可以使用 cuDNN 最佳化的後端在 GPU 裝置上更快地執行 RNN 序列。我們當然不想錯過加速訓練迴圈的機會！要使用它，我們只需告訴 LSTM 模組在損失計算時以“迴圈模式”執行。由於我們通常希望有兩個 LSTM 模組的副本，我們透過呼叫 set_recurrent_mode() 方法來做到這一點，該方法將返回一個新的 LSTM 例項（共享權重），該例項將假定輸入資料具有順序性。

policy = Seq(feature, lstm.set_recurrent_mode(True), mlp, qval)

由於我們還有一些未初始化的引數，在建立最佳化器等之前，我們應該先初始化它們。

policy(env.reset())

DQN 損失

我們的 DQN 損失需要我們傳入策略，並且再次需要動作空間。雖然這看起來是重複的，但很重要，因為我們要確保 DQNLoss 和 QValueModule 類是相容的，但彼此之間沒有強依賴關係。

為了使用 Double-DQN，我們需要一個 delay_value 引數，它將建立一個網路引數的非可微副本作為目標網路。

loss_fn = DQNLoss(policy, action_space=env.action_spec, delay_value=True)

由於我們使用的是雙 DQN，我們需要更新目標引數。我們將使用 SoftUpdate 例項來執行這項工作。

updater = SoftUpdate(loss_fn, eps=0.95)

optim = torch.optim.Adam(policy.parameters(), lr=3e-4)

Collector 和回放緩衝區

我們構建了最簡單的資料 collector。我們將嘗試用一百萬幀訓練演算法，每次擴充套件緩衝區 50 幀。緩衝區將設計用於儲存 2 萬個軌跡，每個軌跡 50 步。在每個最佳化步驟（每次資料收集 16 次）中，我們將從緩衝區中收集 4 個條目，總共 200 個轉換。我們將使用 LazyMemmapStorage 儲存將資料儲存在磁碟上。

注意

為了提高效率，我們在這裡只執行幾千次迭代。在實際設定中，總幀數應設定為 100 萬。

collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200, device=device)
rb = TensorDictReplayBuffer(
    storage=LazyMemmapStorage(20_000), batch_size=4, prefetch=10
)

訓練迴圈

為了跟蹤進度，我們每進行 50 次資料收集就在環境中執行一次策略，並在訓練後繪製結果。

utd = 16
pbar = tqdm.tqdm(total=1_000_000)
longest = 0

traj_lens = []
for i, data in enumerate(collector):
    if i == 0:
        print(
            "Let us print the first batch of data.\nPay 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",
            data,
        )
    pbar.update(data.numel())
    # it is important to pass data that is not flattened
    rb.extend(data.unsqueeze(0).to_tensordict().cpu())
    for _ in range(utd):
        s = rb.sample().to(device, non_blocking=True)
        loss_vals = loss_fn(s)
        loss_vals["loss"].backward()
        optim.step()
        optim.zero_grad()
    longest = max(longest, data["step_count"].max().item())
    pbar.set_description(
        f"steps: {longest}, loss_val: {loss_vals['loss'].item(): 4.4f}, action_spread: {data['action'].sum(0)}"
    )
    exploration_module.step(data.numel())
    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 transform 指示 LSTM 模組需要重置

• 將此模組整合到策略和損失模組中

• 確保 collector 瞭解迴圈狀態條目，以便它們可以與其餘資料一起儲存在回放緩衝區中

延伸閱讀

• TorchRL 文件可以在此處找到。