注意
點選此處下載完整示例程式碼
使用 TorchRL 進行強化學習 (PPO) 教程¶
創建於: 2023 年 3 月 15 日 | 最後更新於: 2025 年 3 月 20 日 | 最後驗證於: 2024 年 11 月 5 日
作者: 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 的改進版本。更多資訊請參見論文 Proximal Policy Optimization Algorithms。
PPO 通常被認為是用於線上、同策略強化學習演算法的快速高效方法。TorchRL 提供了一個損失模組，可以為你完成所有工作，因此你可以依賴此實現並專注於解決問題，而無需在每次想要訓練策略時都重新發明輪子。
為了完整起見，此處簡要概述損失的計算方式，儘管這已由我們的 ClipPPOLoss 模組處理——該演算法工作原理如下：1. 我們將在環境中執行策略，收集給定步數的資料批次。2. 然後，使用 REINFORCE 損失的剪裁版本，對此批次的隨機子樣本執行給定次數的最佳化步驟。3. 剪裁將對我們的損失設定一個悲觀邊界：與較高的回報估計相比，較低的回報估計將受到青睞。精確的損失公式為
該損失包含兩個組成部分：在最小化運算元的第一部分，我們簡單地計算了 REINFORCE 損失的加權重要性版本（例如，我們已經校正了當前策略配置滯後於用於資料收集的策略配置這一事實）。最小化運算元的第二部分是類似的損失，我們在比率超過或低於給定一對閾值時對其進行了剪裁。
該損失確保無論優勢是正還是負，都會阻止導致與先前配置顯著偏移的策略更新。
本教程結構如下
首先，我們將定義一組用於訓練的超引數。
接下來，我們將重點介紹如何使用 TorchRL 的包裝器和變換建立我們的環境或模擬器。
接下來，我們將設計策略網路和值模型，這對於損失函式是必不可少的。這些模組將用於配置我們的損失模組。
接下來，我們將建立回放緩衝區和資料載入器。
最後，我們將執行訓練迴圈並分析結果。
在本教程中，我們將使用 tensordict 庫。TensorDict 是 TorchRL 的通用語言：它幫助我們抽象模組的讀寫內容，從而減少關注具體的資料描述，更多關注演算法本身。
import warnings
warnings.filterwarnings("ignore")
from torch import multiprocessing
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 = 50_000
PPO 引數¶
每次資料收集（或批次收集）時，我們將在一定數量的 epoch 內執行最佳化，每次都在巢狀的訓練迴圈中消耗剛剛獲取的全部資料。此處，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
定義環境¶
在強化學習中，環境通常是我們指代模擬器或控制系統的方式。有各種庫提供強化學習的模擬環境，包括 Gymnasium（以前的 OpenAI Gym）、DeepMind control suite 和許多其他庫。作為一個通用庫，TorchRL 的目標是為大量強化學習模擬器提供可互換的介面，使你能夠輕鬆地用一個環境替換另一個環境。例如，建立包裝後的 gym 環境只需幾個字元即可完成
base_env = GymEnv("InvertedDoublePendulum-v4", device=device)
這段程式碼中有幾點需要注意：首先，我們透過呼叫 GymEnv 包裝器建立了環境。如果傳入額外的關鍵字引數，它們將被傳遞給 gym.make 方法，從而覆蓋最常見的環境構造命令。另外，也可以直接使用 gym.make(env_name, **kwargs) 建立一個 gym 環境，然後將其包裝在 GymWrapper 類中。
還有 device 引數：對於 gym，這隻控制輸入動作和觀察狀態儲存在哪個裝置上，但執行始終在 CPU 上完成。原因很簡單，除非另有說明，否則 gym 不支援裝置端執行。對於其他庫，我們可以控制執行裝置，並且儘可能在儲存和執行後端方面保持一致。
變換¶
我們將向環境新增一些變換，以準備策略所需的資料。在 Gym 中，這通常透過包裝器實現。TorchRL 採用了不同的方法，與 PyTorch 的其他領域庫更相似，透過使用變換來實現。要向環境新增變換，只需將其包裝在 TransformedEnv 例項中，並將變換序列附加到其上。轉換後的環境將繼承被包裝環境的裝置和元資料，並根據其包含的變換序列進行轉換。
歸一化¶
首先要編碼的是歸一化變換。根據經驗法則，最好使資料大致符合單位高斯分佈：為了實現這一點，我們將在環境中執行一定數量的隨機步驟，並計算這些觀測值的摘要統計資訊。
我們將新增另外兩個變換：DoubleToFloat 變換將雙精度條目轉換為單精度數字，以便策略讀取。StepCounter 變換將用於計算環境終止前的步數。我們將使用此指標作為效能的補充衡量標準。
正如我們稍後將看到的，TorchRL 的許多類依賴於 TensorDict 進行通訊。你可以將其視為一個具有額外張量特徵的 Python 字典。實際上，這意味著我們將使用的許多模組需要知道它們將接收到的 tensordict 中要讀取哪些鍵（in_keys）以及要寫入哪些鍵（out_keys）。通常，如果省略 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() 函式執行一個小型的 rollout，並將其輸出與環境規範進行比較。如果沒有引發錯誤，我們可以確信規範已正確定義
check_env_specs(env)
為了好玩，讓我們看看一個簡單的隨機 rollout 是什麼樣的。你可以呼叫 env.rollout(n_steps)，瞭解環境的輸入和輸出是什麼樣的。動作將自動從動作規範域中抽取，因此你無需關心設計隨機取樣器。
通常，在每一步，強化學習環境接收一個動作作為輸入，並輸出一個觀測、一個獎勵和一個完成狀態。觀測可以是複合的，意味著它可以由多個張量組成。這對 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 提供了這種分佈，我們唯一需要關心的是構建一個神經網路，輸出策略工作所需的正確數量的引數（位置，即均值，和比例）
此處唯一增加的難度是將我們的輸出分成兩個相等的部分，並將第二部分對映到嚴格正空間。
我們分三步設計策略
定義一個神經網路
D_obs->2 * D_action。實際上，我們的loc(mu) 和scale(sigma) 都具有維度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.space.low,
"high": env.action_spec.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 提供了一系列 資料收集器類。簡而言之，這些類執行三個操作：重置環境，根據最新觀測計算動作，在環境中執行一步，然後重複最後兩個步驟，直到環境發出停止訊號（或達到完成狀態）。
它們允許你控制每次迭代收集多少幀（透過 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,
)
回放緩衝區¶
回放緩衝區是離策略強化學習演算法的常見組成部分。在同策略環境中，每次收集到一批資料時都會重新填充回放緩衝區，並且其資料會在一定數量的 epoch 中重複使用。
TorchRL 的回放緩衝區使用通用容器 ReplayBuffer 構建，它接受緩衝區的元件作為引數：儲存、寫入器、取樣器以及可能的變換。只有儲存（指示回放緩衝區的容量）是必需的。我們還指定了一個不重複的取樣器，以避免在一個 epoch 中多次取樣同一項。對於 PPO，使用回放緩衝區不是必需的，我們可以簡單地從收集到的批次中取樣子批次，但使用這些類可以輕鬆地以可重複的方式構建內部訓練迴圈。
replay_buffer = ReplayBuffer(
storage=LazyTensorStorage(max_size=frames_per_batch),
sampler=SamplerWithoutReplacement(),
)
損失函式¶
為了方便起見，PPO 損失可以直接從 TorchRL 中匯入，使用 ClipPPOLoss 類。這是使用 PPO 最簡單的方式：它隱藏了 PPO 的數學運算和相關的控制流。
PPO需要計算一些“優勢估計”。簡而言之，優勢是一個反映回報值預期的值，同時處理偏差/方差的權衡。要計算優勢，只需(1)構建利用我們價值運算元的優勢模組，以及(2)在每個epoch之前將每批資料透過它。GAE模組將使用新的tensordict和"advantage"條目更新輸入的"value_target"。 "value_target"是一個無梯度的張量，表示價值網路應該用輸入觀測值表示的經驗值。這兩者都將由ClipPPOLoss使用，以返回策略損失和價值損失。
advantage_module = GAE(
gamma=gamma, lmbda=lmbda, value_network=value_module, average_gae=True, device=device,
)
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()
結果¶
在達到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。
指令碼總執行時間： ( 0 分鐘 0.000 秒)
