注意
轉到末尾 下載完整示例程式碼。
強化學習 (PPO) 使用 TorchRL 教程¶
作者: 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 通常被認為是一種快速高效的線上、同策略強化學習演算法。TorchRL 提供了一個損失模組，為你處理所有工作，這樣你就可以依賴這個實現，專注於解決你的問題，而不是每次訓練策略時都重複造輪子。
為完整起見，這裡簡要概述一下損失函式計算的內容，即使這部分工作由我們的 ClipPPOLoss 模組負責——該演算法工作流程如下：1. 我們透過在環境中執行策略指定步數來取樣一批資料。2. 然後，我們使用 REINFORCE 損失的裁剪版本，從該批資料中隨機抽取子樣本，執行指定數量的最佳化步驟。3. 裁剪操作將對我們的損失設定一個悲觀界限：與更高的回報估計相比，會偏向更低的回報估計。損失函式的精確公式為
損失函式中有兩個組成部分：在 min 運算子的第一部分，我們簡單地計算 REINFORCE 損失的權重重要性版本（例如，我們已經修正了的 REINFORCE 損失，以彌補當前策略配置滯後於用於資料收集的配置的事實）。min 運算子的第二部分是一個類似的損失函式，我們在比率超出或低於給定閾值對時對其進行了裁剪。
這種損失函式確保，無論優勢是正向還是負向，都會阻止那些會產生相對於先前配置顯著變化的策略更新。
本教程結構如下
首先，我們將定義用於訓練的一組超引數。
接下來，我們將重點介紹如何使用 TorchRL 的包裝器（wrappers）和 transforms 建立我們的環境或模擬器。
接下來，我們將設計策略網路和價值模型，這對於損失函式來說是必不可少的。這些模組將用於配置我們的損失模組。
接下來，我們將建立回放緩衝區和資料載入器。
最後，我們將執行訓練迴圈並分析結果。
在整個教程中，我們將使用 tensordict 庫。TensorDict 是 TorchRL 的通用語言（lingua franca）：它幫助我們抽象出模組的讀寫內容，減少對具體資料描述的關注，更多地關注演算法本身。
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 引數來選擇每個批次的大小。我們還將定義允許使用的幀數（例如與模擬器互動的次數）。一般來說，強化學習演算法的目標是在環境互動方面儘快學會解決任務：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 control suite 等等。作為一個通用庫，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¶
我們將向環境中新增一些 transforms，以準備用於策略的資料。在 Gym 中，這通常透過包裝器（wrappers）實現。TorchRL 採取了不同的方法，更類似於其他 PyTorch 領域庫，透過使用 transforms 來實現。要向環境新增 transforms，只需將其包裝在 TransformedEnv 例項中，並向其附加 transforms 序列。轉換後的環境將繼承被包裝環境的裝置和元資料，並根據其包含的 transforms 序列進行轉換。
歸一化¶
首先進行編碼的是一個歸一化 transform。根據經驗，最好讓資料大致符合單位高斯分佈：為此，我們將在環境中執行一定數量的隨機步驟，並計算這些觀測值的彙總統計資料。
我們將附加另外兩個 transforms：DoubleToFloat transform 將雙精度條目轉換為單精度數字，可供策略讀取。StepCounter transform 將用於計算環境終止前的步驟數。我們將使用此度量作為效能的補充度量。
正如我們稍後將看到的，許多 TorchRL 的類依賴於 TensorDict 進行通訊。你可以將其視為具有一些額外張量特性的 Python 字典。實際上，這意味著我們將要使用的許多模組需要被告知要讀取哪些鍵 (in_keys) 以及要寫入哪些鍵 (out_keys) 到它們將接收的 tensordict 中。通常，如果省略 out_keys，則假定 in_keys 條目將原地更新。對於我們的 transforms，我們唯一感興趣的條目被稱為 "observation"，並且我們的 transform 層將被告知只修改此條目。
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 transform 現在已經填充了位置（location）和尺度（scale），這些將用於歸一化資料。
讓我們對彙總統計資料的形狀做一點健全性檢查。
print("normalization constant shape:", env.transform[0].loc.shape)
環境不僅由其模擬器和 transforms 定義，還由一系列元資料定義，這些元資料描述了在其執行期間可以預期什麼。出於效率考慮，TorchRL 對環境規範相當嚴格，但你可以輕鬆檢查你的環境規範是否足夠。在我們的示例中，GymWrapper 和繼承自它的 GymEnv 已經負責為你的環境設定適當的規範，因此你無需為此擔心。
儘管如此，讓我們透過檢視轉換後環境的規範（specs）來看一個具體示例。有三個規範需要檢視：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)
the check_env_specs() 函式執行一個小的 rollout，並將其輸出與環境規範進行比較。如果沒有引發錯誤，我們可以確信規範已正確定義。
check_env_specs(env)
作為一個趣味示例，讓我們看看一個簡單的隨機 rollout 是什麼樣子。你可以呼叫 env.rollout(n_steps) 並獲取環境輸入和輸出的概覽。動作將自動從 action spec 域中抽取，因此你無需關心設計一個隨機取樣器。
通常，在每個步驟中，RL 環境接收動作作為輸入，並輸出觀測值、獎勵和完成狀態。觀測值可以是複合的，意味著它可能由多個張量組成。這對 TorchRL 來說不是問題，因為整組觀測值會自動打包在輸出 TensorDict 中。在給定步數上執行 rollout（例如，一系列環境步驟和隨機動作生成）後，我們將檢索到一個 TensorDict 例項，其形狀與此軌跡長度匹配。
rollout = env.rollout(3)
print("rollout of three steps:", rollout)
print("Shape of the rollout TensorDict:", rollout.batch_size)
我們的 rollout 資料形狀為 torch.Size([3])，與我們執行的步數相匹配。"next" 條目指向當前步驟之後的資料。在大多數情況下，時間 t 的 "next" 資料與時間 t+1 的資料匹配，但如果我們使用某些特定的轉換（例如，多步），則可能不是這種情況。
策略¶
PPO 利用隨機策略來處理探索。這意味著我們的神經網路必須輸出一個分佈的引數，而不是與所採取的動作相對應的單個值。
由於資料是連續的，我們使用 Tanh-Normal 分佈來尊重動作空間的邊界。TorchRL 提供了這種分佈，我們唯一需要關心的是構建一個神經網路，它輸出策略正常工作所需的正確數量的引數（一個位置 loc，即均值，和一個尺度 scale）。
這裡帶來的唯一額外困難是將我們的輸出分成兩個相等的部分，並將第二部分對映到一個嚴格為正的空間。
我們分三步設計策略
定義一個神經網路
D_obs->2 * D_action。實際上，我們的loc(mu) 和scale(sigma) 的維度都為D_action。附加一個
NormalParamExtractor來提取位置（location）和尺度（scale）（例如，將輸入分成兩個相等的部分，並對尺度引數應用一個正向轉換）。建立一個機率型
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()))
資料收集器¶
TorchRL 提供了一組 DataCollector 類。簡而言之，這些類執行三個操作：重置環境，根據最新的觀測計算動作，在環境中執行一步，並重復最後兩個步驟，直到環境發出停止訊號（或達到完成狀態）。
它們允許你控制每次迭代收集多少幀（透過 frames_per_batch 引數），何時重置環境（透過 max_frames_per_traj 引數），策略應在哪個 device 上執行等等。它們還設計用於高效處理批處理和多程序環境。
最簡單的資料收集器是 SyncDataCollector：它是一個迭代器，你可以用它來獲取給定長度的資料批次，並在收集到總幀數 (total_frames) 後停止。其他資料收集器（MultiSyncDataCollector 和 MultiaSyncDataCollector）將在一組多程序工作程式上以同步和非同步方式執行相同的操作。
與之前的策略和環境一樣，資料收集器將返回總元素數與 frames_per_batch 匹配的 TensorDict 例項。使用 TensorDict 將資料傳遞給訓練迴圈，使你能夠編寫完全不關心 rollout 內容具體細節的資料載入管道。
collector = SyncDataCollector(
env,
policy_module,
frames_per_batch=frames_per_batch,
total_frames=total_frames,
split_trajs=False,
device=device,
)
回放緩衝區¶
回放緩衝區是離策略 RL 演算法的常見組成部分。在同策略上下文中，回放緩衝區在每次收集一批資料時都會重新填充，並且其資料會在一定數量的週期（epochs）內被重複使用。
TorchRL 的回放緩衝區是使用通用容器 ReplayBuffer 構建的，該容器接受緩衝區的元件作為引數：一個 storage、一個 writer、一個 sampler 以及可能的 transforms。只有 storage（表示回放緩衝區容量）是必需的。我們還指定了一個無重複取樣器，以避免在一個週期內多次取樣同一項。對於 PPO 來說，使用回放緩衝區並非強制，我們可以簡單地從收集的批次中取樣子批次，但使用這些類可以讓我們以可復現的方式輕鬆構建內部訓練迴圈。
replay_buffer = ReplayBuffer(
storage=LazyTensorStorage(max_size=frames_per_batch),
sampler=SamplerWithoutReplacement(),
)
損失函式¶
為了方便，可以直接從 TorchRL 匯入 PPO 損失函式，使用 ClipPPOLoss 類。這是使用 PPO 最簡單的方式：它隱藏了 PPO 的數學運算及其控制流程。
PPO 需要計算一些“優勢估計”（advantage estimation）。簡而言之，優勢是一個值，它反映了對回報值的期望，同時處理偏置/方差權衡（bias / variance tradeoff）。要計算優勢，只需 (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()
結果¶
在達到 1M 步限制之前，演算法應該已經達到了 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瞭解更多資訊。
