注意

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

使用 TorchRL 訓練強化學習 (PPO) 教程¶

作者：Vincent Moens

本教程演示如何使用 PyTorch 和 torchrl 來訓練引數化策略網路，以解決來自 OpenAI-Gym/Farama-Gymnasium 控制庫的倒立擺任務。

主要學習內容

如何在 TorchRL 中建立環境、轉換其輸出以及從環境中收集資料；
如何使用 TensorDict 讓類之間相互通訊；
使用 TorchRL 構建訓練迴圈的基礎知識
- 如何計算策略梯度方法的優勢訊號；
- 如何使用機率神經網路建立隨機策略；
- 如何建立動態回放緩衝區並從中進行不重複取樣。

我們將介紹 TorchRL 的六個關鍵元件

如果您在 Google Colab 中執行此程式碼，請確保安裝以下依賴項

!pip3 install torchrl
!pip3 install gym[mujoco]
!pip3 install tqdm

近端策略最佳化（PPO）是一種策略梯度演算法，它收集並直接消耗一批資料，以在存在某些近端約束的情況下最大化預期回報來訓練策略。您可以將其視為 REINFORCE（基礎策略最佳化演算法）的複雜版本。有關更多資訊，請參閱近端策略最佳化演算法論文。

PPO 通常被認為是用於線上、on-policy 強化演算法的快速高效方法。TorchRL 提供了一個損失模組，可以為您完成所有工作，因此您可以依賴此實現，專注於解決問題，而不是每次想要訓練策略時都重新發明輪子。

為完整起見，這裡是對損失計算的簡要概述，儘管這由我們的 ClipPPOLoss 模組處理—演算法如下： 1. 我們將透過在環境中執行策略給定的步數來取樣一個數據批。 2. 然後，我們將使用裁剪版的 REINFORCE 損失對該批資料進行隨機子取樣，並執行給定次數的最佳化步驟。 3. 裁剪將對我們的損失設定一個悲觀下界：相比於較高的回報估計，較低的回報估計將受到青睞。損失的確切公式為

\[L(s,a,\theta_k,\theta) = \min\left( \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)} A^{\pi_{\theta_k}}(s,a), \;\; g(\epsilon, A^{\pi_{\theta_k}}(s,a)) \right),\]

該損失有兩個組成部分：在最小運算子的第一部分，我們簡單地計算 REINFORCE 損失的樣本加權版本（例如，一個 REINFORCE 損失，我們已根據當前策略配置滯後於資料收集所用配置的事實進行了校正）。該最小運算子的第二部分是類似的損失，其中我們已將比率裁剪到給定的閾值對之外或之內。

此損失確保無論優勢是正數還是負數，都會抑制會產生與先前配置顯著變化的策略更新。

本教程結構如下

首先，我們將定義一組用於訓練的超引數。
接下來，我們將重點關注使用 TorchRL 的包裝器和變換來建立我們的環境或模擬器。
接下來，我們將設計策略網路和價值模型，這對於損失函式至關重要。這些模組將用於配置我們的損失模組。
接下來，我們將建立回放緩衝區和資料載入器。
最後，我們將執行訓練迴圈並分析結果。

在本教程中，我們將使用 tensordict 庫。 TensorDict 是 TorchRL 的通用語言：它幫助我們抽象模組讀取和寫入的內容，從而減少對特定資料描述的關注，而更多地關注演算法本身。

from collections import defaultdict

import matplotlib.pyplot as plt
import torch
from tensordict.nn import TensorDictModule
from tensordict.nn.distributions import NormalParamExtractor
from torch import nn

from torchrl.collectors import SyncDataCollector
from torchrl.data.replay_buffers import ReplayBuffer
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
from torchrl.data.replay_buffers.storages import LazyTensorStorage
from torchrl.envs import (
    Compose,
    DoubleToFloat,
    ObservationNorm,
    StepCounter,
    TransformedEnv,
)
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.utils import check_env_specs, ExplorationType, set_exploration_type
from torchrl.modules import ProbabilisticActor, TanhNormal, ValueOperator
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value import GAE
from tqdm import tqdm

定義超引數¶

我們為演算法設定超引數。根據可用的資源，可以選擇在 GPU 或其他裝置上執行策略。frame_skip 將控制單個動作執行多少幀。其他計算幀數的引數必須針對此值進行校正（因為一個環境步驟實際上將返回 frame_skip 幀）。

is_fork = multiprocessing.get_start_method() == "fork"
device = (
    torch.device(0)
    if torch.cuda.is_available() and not is_fork
    else torch.device("cpu")
)
num_cells = 256  # number of cells in each layer i.e. output dim.
lr = 3e-4
max_grad_norm = 1.0

資料收集引數¶

在收集資料時，我們可以透過定義 frames_per_batch 引數來選擇每個批次的大小。我們還將定義允許使用的幀數（例如，與模擬器的互動次數）。通常，RL 演算法的目標是學會盡快解決任務（以環境互動的次數衡量）：total_frames 越低越好。

frames_per_batch = 1000
# For a complete training, bring the number of frames up to 1M
total_frames = 10_000

PPO 引數¶

在每次資料收集（或批次收集）時，我們將在一定數量的 *epochs* 上執行最佳化，每次都消耗我們剛剛獲取的整個資料，並在一個巢狀的訓練迴圈中進行。這裡，sub_batch_size 不同於上面的 frames_per_batch：請記住，我們處理的是來自收集器的一個“資料批”，其大小由 frames_per_batch 定義，並且在內部訓練迴圈中，我們將進一步將其分成更小的子批次。這些子批次的大小由 sub_batch_size 控制。

sub_batch_size = 64  # cardinality of the sub-samples gathered from the current data in the inner loop
num_epochs = 10  # optimization steps per batch of data collected
clip_epsilon = (
    0.2  # clip value for PPO loss: see the equation in the intro for more context.
)
gamma = 0.99
lmbda = 0.95
entropy_eps = 1e-4

定義環境¶

在 RL 中，*環境* 通常是我們用來指代模擬器或控制系統的術語。各種庫提供了強化學習的模擬環境，包括 Gymnasium（以前稱為 OpenAI Gym）、DeepMind 控制套件等。作為通用庫，TorchRL 的目標是提供與各種 RL 模擬器可互換的介面，讓您可以輕鬆地將一個環境替換為另一個。例如，使用幾個字元就可以建立包裝好的 gym 環境

base_env = GymEnv("InvertedDoublePendulum-v4", device=device)

此程式碼中有幾點需要注意：首先，我們透過呼叫 GymEnv 包裝器來建立環境。如果傳遞了額外的關鍵字引數，它們將被傳遞給 gym.make 方法，從而涵蓋最常見的環境構造命令。或者，也可以直接使用 gym.make(env_name, **kwargs) 建立一個 gym 環境，並將其包裝在 GymWrapper 類中。

還有 device 引數：對於 gym，這僅控制儲存輸入動作和觀察狀態的裝置，但執行始終在 CPU 上進行。原因很簡單，gym 不支援裝置執行，除非另有說明。對於其他庫，我們可以控制執行裝置，並且在我們能力範圍內，我們努力在儲存和執行後端方面保持一致。

變換 (Transforms)¶

我們將向我們的環境新增一些變換，以準備策略的資料。在 Gym 中，這通常透過包裝器實現。TorchRL 採用不同的方法，更類似於其他 PyTorch 領域庫，透過使用變換。要向環境新增變換，只需將其包裝在 TransformedEnv 例項中，並附加變換序列。轉換後的環境將繼承被包裝環境的裝置和元資料，並根據其包含的變換序列轉換這些內容。

標準化¶

首先要編碼的是標準化變換。經驗法則規定，最好使資料大致匹配單位高斯分佈：為了實現這一點，我們將執行一定數量的隨機步驟，並計算這些觀察的統計資料。

我們將新增另外兩個變換：DoubleToFloat 變換會將雙精度條目轉換為單精度數字，以便策略讀取。StepCounter 變換將用於計算環境終止前的步數。我們將使用此度量作為效能的補充度量。

正如我們稍後將看到的，TorchRL 的許多類都依賴 TensorDict 進行通訊。您可以將其視為具有一些額外張量功能的 Python 字典。實際上，這意味著我們將處理的許多模組都需要被告知讀取哪個鍵（in_keys）以及寫入哪個鍵（out_keys）在它們將收到的 tensordict 中。通常，如果省略 out_keys，則假定 in_keys 條目將被原地更新。對於我們的變換，唯一感興趣的條目是 "observation"，我們將告知變換層修改此條目，僅此條目。

env = TransformedEnv(
    base_env,
    Compose(
        # normalize observations
        ObservationNorm(in_keys=["observation"]),
        DoubleToFloat(),
        StepCounter(),
    ),
)

正如您可能注意到的，我們建立了一個標準化層，但沒有設定其標準化引數。要做到這一點，ObservationNorm 可以自動收集我們環境的統計資料

env.transform[0].init_stats(num_iter=1000, reduce_dim=0, cat_dim=0)

現在，ObservationNorm 變換已填充了位置和縮放因子，將用於標準化資料。

讓我們對統計資料的形狀進行一些基本的檢查

print("normalization constant shape:", env.transform[0].loc.shape)

環境不僅由其模擬器和變換定義，還由一系列元資料定義，這些元資料描述了在執行期間可以預期什麼。為了提高效率，TorchRL 在環境規範方面非常嚴格，但您可以輕鬆檢查您的環境規範是否足夠。在我們的示例中，GymWrapper 和繼承自它的 GymEnv 已經負責為您的環境設定正確的規範，因此您不必關心這一點。

儘管如此，讓我們透過檢視我們轉換後的環境的規範來舉一個具體的例子。有三個規範需要檢視：observation_spec 定義了在環境中執行動作時可以預期什麼，reward_spec 指示了獎勵域，最後是 input_spec（包含 action_spec），它代表了環境執行單個步驟所需的一切。

print("observation_spec:", env.observation_spec)
print("reward_spec:", env.reward_spec)
print("input_spec:", env.input_spec)
print("action_spec (as defined by input_spec):", env.action_spec)

該 check_env_specs() 函式執行一個小的回滾，並將其輸出與環境規範進行比較。如果沒有引發錯誤，我們可以確信規範已正確定義。

check_env_specs(env)

為了好玩，讓我們看看一個簡單的隨機滾動是什麼樣的。您可以呼叫 env.rollout(n_steps) 並大致瞭解環境輸入和輸出的樣子。動作將自動從動作規範域中抽取，因此您不必擔心設計一個隨機取樣器。

通常，在每一步，RL 環境都會接收一個動作作為輸入，並輸出一個觀察、一個獎勵和一個完成狀態。觀察可能是複合的，這意味著它可能由多個張量組成。這對於 TorchRL 來說不是問題，因為整個觀察集會自動打包到輸出的 TensorDict 中。在執行給定步數的滾動（例如，一系列環境步驟和隨機動作生成）後，我們將檢索一個 TensorDict 例項，其形狀與此軌跡長度匹配

rollout = env.rollout(3)
print("rollout of three steps:", rollout)
print("Shape of the rollout TensorDict:", rollout.batch_size)

我們的滾動資料形狀為 torch.Size([3])，與我們執行的步數匹配。"next" 條目指向當前步驟之後的資料。在大多數情況下，時間 t 的 "next" 資料與 t+1 的資料匹配，但如果使用了某些特定變換（例如，多步），則可能不匹配。

策略 (Policy)¶

PPO 使用隨機策略來處理探索。這意味著我們的神經網路將不得不輸出一個分佈的引數，而不是對應於所採取動作的單個值。

由於資料是連續的，我們使用 Tanh-Normal 分佈來尊重動作空間邊界。TorchRL 提供了這種分佈，我們唯一需要關心的是構建一個輸出正確數量引數的神經網路，以便策略可以使用（一個位置，或均值，和一個縮放因子）

\[f_{\theta}(\text{observation}) = \mu_{\theta}(\text{observation}), \sigma^{+}_{\theta}(\text{observation})\]

這裡帶來的唯一額外困難是將我們的輸出分成兩部分，並將第二部分對映到一個嚴格正的空間。

我們分三步設計策略

定義一個神經網路 D_obs -> 2 * D_action。事實上，我們的 loc（均值）和 scale（標準差）都具有 D_action 的維度。
附加一個 NormalParamExtractor 來提取位置和縮放因子（例如，將輸入分成兩等份並將縮放參數應用正變換）。
建立一個機率性的 TensorDictModule，它可以生成此分佈並從中取樣。

actor_net = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(2 * env.action_spec.shape[-1], device=device),
    NormalParamExtractor(),
)

為了讓策略能夠透過 tensordict 資料載體與環境“對話”，我們將 nn.Module 包裝在 TensorDictModule 中。此類將簡單地讀取它提供的 in_keys，並在註冊的 out_keys 處原地寫入輸出。

policy_module = TensorDictModule(
    actor_net, in_keys=["observation"], out_keys=["loc", "scale"]
)

我們現在需要根據正態分佈的位置和縮放因子構建一個分佈。為此，我們指示 ProbabilisticActor 類構建一個 TanhNormal。我們還提供了該分佈的最小值和最大值，這些值是從環境規範中獲取的。

in_keys 的名稱（因此也是上面 TensorDictModule 的 out_keys 的名稱）不能設定為任何您喜歡的字串，因為 TanhNormal 分佈建構函式將期望 loc 和 scale 關鍵字引數。話雖如此，ProbabilisticActor 還接受 Dict[str, str] 型別的 in_keys，其中鍵值對指示應為要使用的每個關鍵字引數使用哪個 in_key 字串。

policy_module = ProbabilisticActor(
    module=policy_module,
    spec=env.action_spec,
    in_keys=["loc", "scale"],
    distribution_class=TanhNormal,
    distribution_kwargs={
        "low": env.action_spec_unbatched.space.low,
        "high": env.action_spec_unbatched.space.high,
    },
    return_log_prob=True,
    # we'll need the log-prob for the numerator of the importance weights
)

價值網路¶

價值網路是 PPO 演算法的關鍵組成部分，儘管它不會在推理時使用。此模組將讀取觀察值並返回對後續軌跡的折扣回報的估計。這使我們能夠透過依賴在訓練過程中即時學習的某種效用估計來攤銷學習。我們的價值網路與策略具有相同的結構，但為了簡單起見，我們為其分配了自己的引數集。

value_net = nn.Sequential(
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(num_cells, device=device),
    nn.Tanh(),
    nn.LazyLinear(1, device=device),
)

value_module = ValueOperator(
    module=value_net,
    in_keys=["observation"],
)

讓我們嘗試我們的策略和價值模組。如前所述，TensorDictModule 的使用使得可以直接讀取環境的輸出來執行這些模組，因為它們知道要讀取哪些資訊以及寫入何處

print("Running policy:", policy_module(env.reset()))
print("Running value:", value_module(env.reset()))

資料收集器 (Data collector)¶

TorchRL 提供了一組資料收集器類。簡而言之，這些類執行三個操作：重置環境，根據最新觀察計算動作，在環境中執行一步，然後重複最後兩個步驟，直到環境發出停止訊號（或達到完成狀態）。

它們允許您控制每次迭代收集多少幀（透過 frames_per_batch 引數），何時重置環境（透過 max_frames_per_traj 引數），策略應在哪個 device 上執行，等等。它們還旨在與批次和多程序環境高效地協同工作。

最簡單的資料收集器是 SyncDataCollector：它是一個迭代器，您可以使用它來獲取給定長度的資料批，並在收集完總幀數（total_frames）後停止。其他資料收集器（MultiSyncDataCollector 和 MultiaSyncDataCollector）將在多程序工作者集合上以同步和非同步方式執行相同的操作。

與之前的策略和環境一樣，資料收集器將返回 TensorDict 例項，其總元素數量將匹配 frames_per_batch。使用 TensorDict 將資料傳遞給訓練迴圈允許您編寫 100% 忽略滾動內容實際特異性的資料載入管道。

collector = SyncDataCollector(
    env,
    policy_module,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device,
)

回放緩衝區 (Replay buffer)¶

回放緩衝區是離策略 RL 演算法的常見構建模組。在策略環境中，每當收集一批資料時，回放緩衝區就會被重新填充，並且其資料會在一定數量的 epoch 中被重複消耗。

TorchRL 的回放緩衝區是使用通用容器 ReplayBuffer 構建的，它接受緩衝區的元件作為引數：儲存、寫入器、取樣器以及可能的變換。只有儲存（指示回放緩衝區容量）是必需的。我們還指定了一個無重複取樣器，以避免在一個 epoch 中多次取樣同一項。將回放緩衝區用於 PPO 不是必需的，我們可以簡單地從收集的資料批中取樣子批次，但使用這些類可以輕鬆地以可重現的方式構建內部訓練迴圈。

replay_buffer = ReplayBuffer(
    storage=LazyTensorStorage(max_size=frames_per_batch),
    sampler=SamplerWithoutReplacement(),
)

損失函式 (Loss function)¶

PPO 損失可以方便地從 TorchRL 中直接匯入，使用 ClipPPOLoss 類。這是使用 PPO 的最簡單方法：它隱藏了 PPO 的數學運算以及與之相關的控制流。

PPO 需要計算一些“優勢估計”。簡而言之，優勢是一個值，它反映了在處理偏差/方差權衡時的期望回報值。要計算優勢，只需（1）構建優勢模組，該模組利用我們的價值運算子，以及（2）在每個 epoch 之前將每個資料批透過它。GAE 模組將使用新的 "advantage" 和 "value_target" 條目更新輸入的 tensordict。"value_target" 是一個無梯度張量，代表了價值網路應與輸入觀察值關聯的經驗值。兩者都將被 ClipPPOLoss 用於返回策略和價值損失。

advantage_module = GAE(
    gamma=gamma, lmbda=lmbda, value_network=value_module, average_gae=True
)

loss_module = ClipPPOLoss(
    actor_network=policy_module,
    critic_network=value_module,
    clip_epsilon=clip_epsilon,
    entropy_bonus=bool(entropy_eps),
    entropy_coef=entropy_eps,
    # these keys match by default but we set this for completeness
    critic_coef=1.0,
    loss_critic_type="smooth_l1",
)

optim = torch.optim.Adam(loss_module.parameters(), lr)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optim, total_frames // frames_per_batch, 0.0
)

訓練迴圈¶

現在我們有了編寫訓練迴圈所需的所有元件。步驟包括

收集資料
- 計算優勢
  - 迴圈遍歷收集的資料以計算損失值
  - 反向傳播
  - 最佳化
  - 重複
- 重複
重複

logs = defaultdict(list)
pbar = tqdm(total=total_frames)
eval_str = ""

# We iterate over the collector until it reaches the total number of frames it was
# designed to collect:
for i, tensordict_data in enumerate(collector):
    # we now have a batch of data to work with. Let's learn something from it.
    for _ in range(num_epochs):
        # We'll need an "advantage" signal to make PPO work.
        # We re-compute it at each epoch as its value depends on the value
        # network which is updated in the inner loop.
        advantage_module(tensordict_data)
        data_view = tensordict_data.reshape(-1)
        replay_buffer.extend(data_view.cpu())
        for _ in range(frames_per_batch // sub_batch_size):
            subdata = replay_buffer.sample(sub_batch_size)
            loss_vals = loss_module(subdata.to(device))
            loss_value = (
                loss_vals["loss_objective"]
                + loss_vals["loss_critic"]
                + loss_vals["loss_entropy"]
            )

            # Optimization: backward, grad clipping and optimization step
            loss_value.backward()
            # this is not strictly mandatory but it's good practice to keep
            # your gradient norm bounded
            torch.nn.utils.clip_grad_norm_(loss_module.parameters(), max_grad_norm)
            optim.step()
            optim.zero_grad()

    logs["reward"].append(tensordict_data["next", "reward"].mean().item())
    pbar.update(tensordict_data.numel())
    cum_reward_str = (
        f"average reward={logs['reward'][-1]: 4.4f} (init={logs['reward'][0]: 4.4f})"
    )
    logs["step_count"].append(tensordict_data["step_count"].max().item())
    stepcount_str = f"step count (max): {logs['step_count'][-1]}"
    logs["lr"].append(optim.param_groups[0]["lr"])
    lr_str = f"lr policy: {logs['lr'][-1]: 4.4f}"
    if i % 10 == 0:
        # We evaluate the policy once every 10 batches of data.
        # Evaluation is rather simple: execute the policy without exploration
        # (take the expected value of the action distribution) for a given
        # number of steps (1000, which is our ``env`` horizon).
        # The ``rollout`` method of the ``env`` can take a policy as argument:
        # it will then execute this policy at each step.
        with set_exploration_type(ExplorationType.DETERMINISTIC), torch.no_grad():
            # execute a rollout with the trained policy
            eval_rollout = env.rollout(1000, policy_module)
            logs["eval reward"].append(eval_rollout["next", "reward"].mean().item())
            logs["eval reward (sum)"].append(
                eval_rollout["next", "reward"].sum().item()
            )
            logs["eval step_count"].append(eval_rollout["step_count"].max().item())
            eval_str = (
                f"eval cumulative reward: {logs['eval reward (sum)'][-1]: 4.4f} "
                f"(init: {logs['eval reward (sum)'][0]: 4.4f}), "
                f"eval step-count: {logs['eval step_count'][-1]}"
            )
            del eval_rollout
    pbar.set_description(", ".join([eval_str, cum_reward_str, stepcount_str, lr_str]))

    # We're also using a learning rate scheduler. Like the gradient clipping,
    # this is a nice-to-have but nothing necessary for PPO to work.
    scheduler.step()

結果 (Results)¶

在達到 100 萬步上限之前，演算法應該已經達到了 1000 步的最大步數限制，這是軌跡被截斷之前的最大步數。

plt.figure(figsize=(10, 10))
plt.subplot(2, 2, 1)
plt.plot(logs["reward"])
plt.title("training rewards (average)")
plt.subplot(2, 2, 2)
plt.plot(logs["step_count"])
plt.title("Max step count (training)")
plt.subplot(2, 2, 3)
plt.plot(logs["eval reward (sum)"])
plt.title("Return (test)")
plt.subplot(2, 2, 4)
plt.plot(logs["eval step_count"])
plt.title("Max step count (test)")
plt.show()

結論和後續步驟¶

在本教程中，我們學習了

如何使用 torchrl 建立和自定義環境；
如何編寫模型和損失函式；
如何設定典型的訓練迴圈。

如果您想進一步嘗試本教程，可以進行以下修改

從效率角度來看，我們可以並行執行多個模擬以加快資料收集。有關更多資訊，請參閱 ParallelEnv。
從日誌記錄角度來看，可以將 torchrl.record.VideoRecorder 變換新增到環境後，請求渲染以獲得倒立擺執行的視覺渲染。有關更多資訊，請參閱 torchrl.record。

由 Sphinx-Gallery 生成的畫廊