評價此頁
★ ★ ★ ★ ★
intermediate/dqn_with_rnn_tutorial

注意

跳轉到末尾 下載完整的示例程式碼。

迴圈 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 中。

Data collection with a recurrent policy

正如該圖所示，我們的環境使用零初始化的迴圈狀態填充 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 環境的例項，並進行一些自定義轉換：轉換為灰度、調整大小到 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()

策略#

我們的策略將包含 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 類來將 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 模組類，特別是它的 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())

就是這樣！現在我們可以列印環境，檢查在新增 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(spec=env.action_spec)

注意

TorchRL 還提供了一個包裝器類 torchrl.modules.QValueActor，它將一個模組包裝在 Sequential 中，並附帶一個 QValueModule，就像我們在這裡顯式做的那樣。這樣做的好處很小，而且過程不太透明，但最終結果將與我們這裡所做的相似。

我們現在可以將它們組合到一個 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 模組時告訴它以“recurrent-mode”執行。由於我們通常會想要兩個 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)

由於我們正在使用 Double DQN，因此我們需要更新目標引數。我們將使用 SoftUpdate 例項來完成此工作。

updater = SoftUpdate(loss_fn, eps=0.95)

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

收集器和回放緩衝區#

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

注意

為了效率，我們這裡只運行了幾千次迭代。在實際設定中，總幀數應設定為 1M。

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 轉換指示 LSTM 模組需要重置

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

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

進一步閱讀#

  • TorchRL 文件可以在 這裡 找到。

© 版權所有 2024, PyTorch。

使用 PyData Sphinx 主題 0.15.4 構建。

在此頁面上
PyTorch 庫