注意
前往結尾以下載完整的範例程式碼。
使用 TorchRL 的強化學習 (PPO) 教學¶
作者: Vincent Moens
本教學示範如何使用 PyTorch 和 torchrl 訓練參數化策略網路，以解決來自 OpenAI-Gym/Farama-Gymnasium 控制函式庫的 Inverted Pendulum 任務。
倒立擺¶
主要學習內容
如何在 TorchRL 中建立環境、轉換其輸出，以及從此環境收集資料；
如何使用
TensorDict使您的類別彼此溝通；使用 TorchRL 建構訓練迴圈的基礎知識
如何計算策略梯度方法的優勢訊號；
如何使用機率神經網路建立隨機策略；
如何建立動態重播緩衝區並從中取樣，而不會重複。
我們將涵蓋 TorchRL 的六個關鍵元件
如果您在 Google Colab 中執行此操作，請確保您安裝以下相依性
!pip3 install torchrl
!pip3 install gym[mujoco]
!pip3 install tqdm
近端策略最佳化 (PPO) 是一種策略梯度演算法，其中收集一批資料並直接使用它來訓練策略，以最大化給定一些近似性約束的預期回報。您可以將其視為 REINFORCE 的複雜版本，REINFORCE 是基礎策略最佳化演算法。有關更多資訊，請參閱 近端策略最佳化演算法論文。
PPO 通常被認為是一種快速有效的線上、基於策略的強化演算法。TorchRL 提供了一個損失模組，可以為您完成所有工作，因此您可以依賴此實作，並專注於解決您的問題，而不是每次想要訓練策略時都重新發明輪子。
為了完整起見，這裡簡要概述了損失的計算方式，即使這由我們的 ClipPPOLoss 模組處理 — 演算法的工作方式如下：1. 我們將透過在環境中播放策略給定的步數來採樣一批資料。2. 然後，我們將使用 REINFORCE 損失的裁剪版本，對此批次執行給定數量的最佳化步驟，並使用隨機子樣本。3. 裁剪將對我們的損失設定一個悲觀的界限：與較高的估計相比，較低的回報估計將會受到青睞。損失的精確公式為
該損失函數包含兩個部分：在最小值運算子的第一部分，我們簡單地計算 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 參數¶
在每次資料收集 (或批次收集) 時，我們將在特定數量的週期上執行最佳化，每次都消耗我們剛在巢狀訓練迴圈中獲得的完整資料。這裡的 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 環境
在此程式碼中有幾件事需要注意：首先，我們透過呼叫 GymEnv 包裝器來建立環境。如果傳遞額外的關鍵字引數，它們將被傳輸到 gym.make 方法，因此涵蓋了最常見的環境建構命令。或者，也可以直接使用 gym.make(env_name, **kwargs) 建立 gym 環境，並將其包裝在 GymWrapper 類別中。
還有 device 引數：對於 gym，這僅控制輸入動作和觀察到的狀態將儲存的裝置，但執行始終在 CPU 上完成。這樣做的原因是 gym 不支援裝置上的執行，除非另有說明。對於其他函式庫，我們可以控制執行裝置，並且盡可能地在儲存和執行後端方面保持一致。
轉換¶
我們將在我們的環境中附加一些轉換，以準備策略的資料。在 Gym 中，這通常透過包裝器來實現。TorchRL 採用了不同的方法，更類似於其他 pytorch 網域函式庫，透過使用轉換。若要將轉換新增到環境，應簡單地將其包裝在 TransformedEnv 實例中，並將轉換序列附加到其中。轉換後的環境將繼承包裝環境的裝置和元資料，並根據其包含的轉換序列來轉換這些資料。
正規化¶
首先要編碼的是正規化轉換。作為經驗法則，最好擁有與單位高斯分佈大致匹配的資料：若要獲得此結果，我們將在環境中執行一定數量的隨機步驟，並計算這些觀察結果的摘要統計資料。
我們將附加另外兩個轉換：DoubleToFloat 轉換將把 double 條目轉換為單精度數字，以便由策略讀取。StepCounter 轉換將用於計算環境終止之前的步驟數。我們將使用此測量作為效能的補充測量。
正如我們稍後會看到的，許多 TorchRL 的類別都依賴 TensorDict 來進行溝通。您可以將其視為一個具有額外張量功能的 Python 字典。實際上，這意味著我們將使用的許多模組都需要被告知從哪個鍵讀取資料 (in_keys) 以及在它們接收到的 tensordict 中將資料寫入哪個鍵 (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)
normalization constant shape: torch.Size([11])
一個環境不僅僅由它的模擬器和轉換來定義，還由一系列描述執行期間預期行為的中繼資料來定義。為了效率起見，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)
observation_spec: Composite(
observation: UnboundedContinuous(
shape=torch.Size([11]),
space=ContinuousBox(
low=Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, contiguous=True),
high=Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, contiguous=True)),
device=cpu,
dtype=torch.float32,
domain=continuous),
step_count: BoundedDiscrete(
shape=torch.Size([1]),
space=ContinuousBox(
low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True),
high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True)),
device=cpu,
dtype=torch.int64,
domain=discrete),
device=cpu,
shape=torch.Size([]))
reward_spec: UnboundedContinuous(
shape=torch.Size([1]),
space=ContinuousBox(
low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
device=cpu,
dtype=torch.float32,
domain=continuous)
input_spec: Composite(
full_state_spec: Composite(
step_count: BoundedDiscrete(
shape=torch.Size([1]),
space=ContinuousBox(
low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True),
high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True)),
device=cpu,
dtype=torch.int64,
domain=discrete),
device=cpu,
shape=torch.Size([])),
full_action_spec: Composite(
action: BoundedContinuous(
shape=torch.Size([1]),
space=ContinuousBox(
low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
device=cpu,
dtype=torch.float32,
domain=continuous),
device=cpu,
shape=torch.Size([])),
device=cpu,
shape=torch.Size([]))
action_spec (as defined by input_spec): BoundedContinuous(
shape=torch.Size([1]),
space=ContinuousBox(
low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
device=cpu,
dtype=torch.float32,
domain=continuous)
check_env_specs() 函數運行一個小的 rollout，並將其輸出與環境規範進行比較。如果沒有引發錯誤，我們可以確信規範已正確定義。
check_env_specs(env)
為了好玩，讓我們看看一個簡單的隨機 rollout 是什麼樣的。您可以調用 env.rollout(n_steps)，並概述環境輸入和輸出的樣子。動作將自動從動作規範域中提取，因此您無需擔心設計隨機採樣器。
通常，在每個步驟中，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 of three steps: TensorDict(
fields={
action: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
next: TensorDict(
fields={
done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
observation: Tensor(shape=torch.Size([3, 11]), device=cpu, dtype=torch.float32, is_shared=False),
reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
step_count: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.int64, is_shared=False),
terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([3]),
device=cpu,
is_shared=False),
observation: Tensor(shape=torch.Size([3, 11]), device=cpu, dtype=torch.float32, is_shared=False),
step_count: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.int64, is_shared=False),
terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([3]),
device=cpu,
is_shared=False)
Shape of the rollout TensorDict: torch.Size([3])
我們的 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()))
Running policy: TensorDict(
fields={
action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
loc: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
observation: Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
sample_log_prob: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
scale: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
step_count: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False),
terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([]),
device=cpu,
is_shared=False)
Running value: TensorDict(
fields={
done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
observation: Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
state_value: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
step_count: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False),
terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([]),
device=cpu,
is_shared=False)
資料收集器¶
TorchRL 提供了一組 DataCollector 類別。簡而言之，這些類別執行三個操作：重置環境、根據最新的觀測值計算動作、在環境中執行一個步驟，並重複最後兩個步驟，直到環境發出停止訊號（或達到完成狀態）。
它們允許您控制每次迭代要收集多少幀（透過 frames_per_batch 參數）、何時重置環境（透過 max_frames_per_traj 參數），以及在哪個 device 上執行策略等等。它們還被設計為可以有效率地與批次處理和多進程環境一起工作。
最簡單的資料收集器是 SyncDataCollector：它是一個迭代器，您可以使用它來取得給定長度的資料批次，並且它會在收集到總幀數 (total_frames) 後停止。其他資料收集器（MultiSyncDataCollector 和 MultiaSyncDataCollector）將在一組多進程工作者上以同步和非同步方式執行相同的操作。
至於之前的策略和環境，資料收集器將回傳 TensorDict 實例，其元素總數將與 frames_per_batch 相符。使用 TensorDict 將資料傳遞到訓練迴圈可讓您編寫 100% 不知道 rollout 內容實際細節的資料載入流程。
collector = SyncDataCollector(
env,
policy_module,
frames_per_batch=frames_per_batch,
total_frames=total_frames,
split_trajs=False,
device=device,
)
重播緩衝區¶
重播緩衝區是離策略 RL 演算法的常見構建模組。在 on-policy 上下文中，每次收集一批資料時都會重新填滿重播緩衝區，並且其資料會重複使用一段時間。
TorchRL 的重播緩衝區是使用一個通用容器 ReplayBuffer 建構的，該容器將緩衝區的元件作為參數：一個儲存、一個寫入器、一個取樣器，可能還有一些轉換。只有儲存（指示重播緩衝區容量）是必需的。我們還指定了一個沒有重複的取樣器，以避免在一個 epoch 中多次取樣相同的項目。為 PPO 使用重播緩衝區不是強制性的，我們可以簡單地從收集的批次中取樣子批次，但使用這些類別可以讓我們以可重複的方式輕鬆地建構內部訓練迴圈。
replay_buffer = ReplayBuffer(
storage=LazyTensorStorage(max_size=frames_per_batch),
sampler=SamplerWithoutReplacement(),
)
損失函數¶
為了方便起見，可以使用 ClipPPOLoss 類別直接從 TorchRL 匯入 PPO 損失。這是利用 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
)
訓練迴圈¶
現在我們擁有編寫訓練迴圈所需的所有部分。步驟包括：
收集資料
計算優勢 (advantage)
循環處理收集到的資料以計算損失值
反向傳播
最佳化
重複
重複
重複
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()
0%| | 0/10000 [00:00<?, ?it/s]
10%|█ | 1000/10000 [00:02<00:19, 460.72it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.0871 (init= 9.0871), step count (max): 12, lr policy: 0.0003: 10%|█ | 1000/10000 [00:02<00:19, 460.72it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.0871 (init= 9.0871), step count (max): 12, lr policy: 0.0003: 20%|██ | 2000/10000 [00:04<00:17, 458.95it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.1306 (init= 9.0871), step count (max): 15, lr policy: 0.0003: 20%|██ | 2000/10000 [00:04<00:17, 458.95it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.1306 (init= 9.0871), step count (max): 15, lr policy: 0.0003: 30%|███ | 3000/10000 [00:06<00:15, 463.69it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.1613 (init= 9.0871), step count (max): 18, lr policy: 0.0003: 30%|███ | 3000/10000 [00:06<00:15, 463.69it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.1613 (init= 9.0871), step count (max): 18, lr policy: 0.0003: 40%|████ | 4000/10000 [00:08<00:12, 465.62it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.1849 (init= 9.0871), step count (max): 20, lr policy: 0.0002: 40%|████ | 4000/10000 [00:08<00:12, 465.62it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.1849 (init= 9.0871), step count (max): 20, lr policy: 0.0002: 50%|█████ | 5000/10000 [00:10<00:10, 467.03it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2013 (init= 9.0871), step count (max): 25, lr policy: 0.0002: 50%|█████ | 5000/10000 [00:10<00:10, 467.03it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2013 (init= 9.0871), step count (max): 25, lr policy: 0.0002: 60%|██████ | 6000/10000 [00:12<00:08, 467.79it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2178 (init= 9.0871), step count (max): 27, lr policy: 0.0001: 60%|██████ | 6000/10000 [00:12<00:08, 467.79it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2178 (init= 9.0871), step count (max): 27, lr policy: 0.0001: 70%|███████ | 7000/10000 [00:14<00:06, 470.65it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2329 (init= 9.0871), step count (max): 32, lr policy: 0.0001: 70%|███████ | 7000/10000 [00:14<00:06, 470.65it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2329 (init= 9.0871), step count (max): 32, lr policy: 0.0001: 80%|████████ | 8000/10000 [00:17<00:04, 458.50it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2473 (init= 9.0871), step count (max): 40, lr policy: 0.0001: 80%|████████ | 8000/10000 [00:17<00:04, 458.50it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2473 (init= 9.0871), step count (max): 40, lr policy: 0.0001: 90%|█████████ | 9000/10000 [00:19<00:02, 463.76it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2477 (init= 9.0871), step count (max): 52, lr policy: 0.0000: 90%|█████████ | 9000/10000 [00:19<00:02, 463.76it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2477 (init= 9.0871), step count (max): 52, lr policy: 0.0000: 100%|██████████| 10000/10000 [00:21<00:00, 467.36it/s]
eval cumulative reward: 91.8959 (init: 91.8959), eval step-count: 9, average reward= 9.2501 (init= 9.0871), step count (max): 38, lr policy: 0.0000: 100%|██████████| 10000/10000 [00:21<00:00, 467.36it/s]
結果¶
在達到 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以了解更多資訊。
腳本的總運行時間： (1 分鐘 26.279 秒)
估計記憶體使用量： 320 MB
