注意
請 轉到末尾 下載完整的示例程式碼。
TorchRL 訓練器：DQN 示例¶
TorchRL 提供了一個通用的 Trainer 類來處理您的訓練迴圈。訓練器執行一個巢狀迴圈，其中外層迴圈是資料收集，內層迴圈消耗這些資料或從回放緩衝區檢索的資料來訓練模型。在訓練迴圈的各個點，可以附加鉤子並在給定時間間隔執行。
在本教程中，我們將使用訓練器類從頭開始訓練一個 DQN 演算法來解決 CartPole 任務。
主要收穫
構建一個包含其基本元件的訓練器：資料收集器、損失模組、回放緩衝區和最佳化器。
向訓練器新增鉤子，例如日誌記錄器、目標網路更新器等。
訓練器是完全可定製的，並提供大量功能。本教程圍繞其構建進行組織。我們將首先詳細介紹如何構建庫的每個元件，然後使用 Trainer 類將這些元件組合在一起。
在此過程中，我們還將關注該庫的一些其他方面
如何在 TorchRL 中構建環境，包括變換（例如，資料歸一化、幀堆疊、調整大小和灰度化）以及並行執行。與我們在 DDPG 教程 中所做的不同，我們將歸一化畫素而不是狀態向量。
如何設計一個
QValueActor物件，即一個估計動作值並選擇估計回報最高的動作的 actor；如何從環境中高效收集資料並將其儲存在回放緩衝區中；
如何使用多步，這是用於離策略演算法的一個簡單預處理步驟；
最後，如何評估您的模型。
先決條件：我們鼓勵您首先透過 PPO 教程 來熟悉 torchrl。
DQN¶
DQN（深度 Q-Learning）是深度強化學習的開創性工作。
從高層次來看，該演算法非常簡單：Q-Learning 包括學習一個狀態-動作值表，以便在遇到任何特定狀態時，我們只需查詢值最高的動作即可知道應選擇哪個動作。這種簡單的設定要求動作和狀態是離散的，否則無法構建查詢表。
DQN 使用一個神經網路，該網路將狀態-動作空間對映到值（標量）空間，從而分攤儲存和探索所有可能的狀態-動作組合的成本：如果過去未見過某個狀態，我們仍然可以將其與透過神經網路的各種可用動作結合起來，並獲得每個可用動作的插值。值。
我們將解決經典的 CartPole 控制問題。摘自該環境檢索的 Gymnasium 文件
我們不旨在提供該演算法的 SOTA 實現，而是為了在演算法的上下文中提供 TorchRL 功能的高階說明。
import os
import uuid
import torch
from torch import nn
from torchrl.collectors import MultiaSyncDataCollector, SyncDataCollector
from torchrl.data import LazyMemmapStorage, MultiStep, TensorDictReplayBuffer
from torchrl.envs import (
EnvCreator,
ExplorationType,
ParallelEnv,
RewardScaling,
StepCounter,
)
from torchrl.envs.libs.gym import GymEnv
from torchrl.envs.transforms import (
CatFrames,
Compose,
GrayScale,
ObservationNorm,
Resize,
ToTensorImage,
TransformedEnv,
)
from torchrl.modules import DuelingCnnDQNet, EGreedyModule, QValueActor
from torchrl.objectives import DQNLoss, SoftUpdate
from torchrl.record.loggers.csv import CSVLogger
from torchrl.trainers import (
LogScalar,
LogValidationReward,
ReplayBufferTrainer,
Trainer,
UpdateWeights,
)
def is_notebook() -> bool:
try:
shell = get_ipython().__class__.__name__
if shell == "ZMQInteractiveShell":
return True # Jupyter notebook or qtconsole
elif shell == "TerminalInteractiveShell":
return False # Terminal running IPython
else:
return False # Other type (?)
except NameError:
return False # Probably standard Python interpreter
讓我們開始處理我們演算法所需的各種元件
一個環境；
一個策略（以及我們將其分組在“模型”下的相關模組）；
一個數據收集器，它使策略在環境中執行並提供訓練資料；
一個回放緩衝區來儲存訓練資料；
一個損失模組，用於計算策略最大化回報的目標函式；
一個最佳化器，它根據我們的損失執行引數更新。
附加模組包括日誌記錄器、記錄器（以“eval”模式執行策略）和目標網路更新器。有了所有這些元件，很容易看出在訓練指令碼中可能會放錯位置或誤用某個元件。訓練器就是為您協調一切！
構建環境¶
首先，讓我們編寫一個將輸出環境的輔助函式。照常，“原始”環境可能過於簡單，無法在實踐中使用，我們需要一些資料轉換來將其輸出暴露給策略。
我們將使用五個變換
StepCounter用於計算每個軌跡中的步數；ToTensorImage將把[W, H, C]的 uint8 張量轉換為[0, 1]空間中的浮點張量，形狀為[C, W, H]；RewardScaling用於減小回報的尺度；GrayScale將把我們的影像轉換為灰度；Resize將把影像調整為 64x64 格式；CatFrames將把任意數量的連續幀（N=4）沿通道維度連線到單個張量中。這很有用，因為單個影像不包含有關 CartPole 運動的資訊。需要一些關於過去觀察和動作的記憶，透過迴圈神經網路或使用幀堆疊。ObservationNorm，它將根據一些自定義摘要統計資訊對我們的觀察進行歸一化。
實際上，我們的環境構建器有兩個引數
parallel：確定是否需要並行執行多個環境。我們將變換堆疊在ParallelEnv之後，以利用裝置上操作的向量化，儘管這在技術上也可以與附加到其自身變換集中的每個單獨環境一起工作。obs_norm_sd將包含ObservationNorm變換的歸一化常數。
def make_env(
parallel=False,
obs_norm_sd=None,
num_workers=1,
):
if obs_norm_sd is None:
obs_norm_sd = {"standard_normal": True}
if parallel:
def maker():
return GymEnv(
"CartPole-v1",
from_pixels=True,
pixels_only=True,
device=device,
)
base_env = ParallelEnv(
num_workers,
EnvCreator(maker),
# Don't create a sub-process if we have only one worker
serial_for_single=True,
mp_start_method=mp_context,
)
else:
base_env = GymEnv(
"CartPole-v1",
from_pixels=True,
pixels_only=True,
device=device,
)
env = TransformedEnv(
base_env,
Compose(
StepCounter(), # to count the steps of each trajectory
ToTensorImage(),
RewardScaling(loc=0.0, scale=0.1),
GrayScale(),
Resize(64, 64),
CatFrames(4, in_keys=["pixels"], dim=-3),
ObservationNorm(in_keys=["pixels"], **obs_norm_sd),
),
)
return env
計算歸一化常數¶
要歸一化影像，我們不希望使用完整的 [C, W, H] 歸一化掩碼獨立地歸一化每個畫素，而是使用更簡單的 [C, 1, 1] 形狀的歸一化常數（位置和尺度引數）集。我們將使用 init_stats() 的 reduce_dim 引數來指示必須約簡的維度，以及 keep_dims 引數以確保並非所有維度都在過程中消失
def get_norm_stats():
test_env = make_env()
test_env.transform[-1].init_stats(
num_iter=1000, cat_dim=0, reduce_dim=[-1, -2, -4], keep_dims=(-1, -2)
)
obs_norm_sd = test_env.transform[-1].state_dict()
# let's check that normalizing constants have a size of ``[C, 1, 1]`` where
# ``C=4`` (because of :class:`~torchrl.envs.CatFrames`).
print("state dict of the observation norm:", obs_norm_sd)
test_env.close()
del test_env
return obs_norm_sd
構建模型（深度 Q 網路）¶
以下函式構建了一個 DuelingCnnDQNet 物件，它是一個簡單的 CNN，後跟一個兩層 MLP。這裡使用的唯一技巧是，動作值（即左和右動作值）是使用以下公式計算的
其中 \(\mathbb{v}\) 是我們的動作值向量，\(b\) 是一個 \(\mathbb{R}^n \rightarrow 1\) 函式，\(v\) 是一個 \(\mathbb{R}^n \rightarrow \mathbb{R}^m\) 函式，對於 \(n = \# obs\) 和 \(m = \# actions\)。
我們的網路被包裝在一個 QValueActor 中，它將讀取狀態-動作值，選擇具有最大值的那個，並將所有這些結果寫入輸入的 tensordict.TensorDict。
def make_model(dummy_env):
cnn_kwargs = {
"num_cells": [32, 64, 64],
"kernel_sizes": [6, 4, 3],
"strides": [2, 2, 1],
"activation_class": nn.ELU,
# This can be used to reduce the size of the last layer of the CNN
# "squeeze_output": True,
# "aggregator_class": nn.AdaptiveAvgPool2d,
# "aggregator_kwargs": {"output_size": (1, 1)},
}
mlp_kwargs = {
"depth": 2,
"num_cells": [
64,
64,
],
"activation_class": nn.ELU,
}
net = DuelingCnnDQNet(
dummy_env.action_spec.shape[-1], 1, cnn_kwargs, mlp_kwargs
).to(device)
net.value[-1].bias.data.fill_(init_bias)
actor = QValueActor(net, in_keys=["pixels"], spec=dummy_env.action_spec).to(device)
# init actor: because the model is composed of lazy conv/linear layers,
# we must pass a fake batch of data through it to instantiate them.
tensordict = dummy_env.fake_tensordict()
actor(tensordict)
# we join our actor with an EGreedyModule for data collection
exploration_module = EGreedyModule(
spec=dummy_env.action_spec,
annealing_num_steps=total_frames,
eps_init=eps_greedy_val,
eps_end=eps_greedy_val_env,
)
actor_explore = TensorDictSequential(actor, exploration_module)
return actor, actor_explore
收集和儲存資料¶
回放緩衝區¶
回放緩衝區在 DQN 等離策略 RL 演算法中起著核心作用。它們構成了訓練期間我們將從中取樣的資料集。
在這裡，我們將使用常規取樣策略，儘管優先回放緩衝區（prioritized RB）可以顯著提高效能。
我們使用 LazyMemmapStorage 類將儲存放在磁碟上。此儲存以惰性方式建立：它將在第一個資料批次傳遞給它後才會被例項化。
此儲存的唯一要求是，在寫入時傳遞給它的資料必須始終具有相同的形狀。
buffer_scratch_dir = tempfile.TemporaryDirectory().name
def get_replay_buffer(buffer_size, n_optim, batch_size, device):
replay_buffer = TensorDictReplayBuffer(
batch_size=batch_size,
storage=LazyMemmapStorage(buffer_size, scratch_dir=buffer_scratch_dir),
prefetch=n_optim,
transform=lambda td: td.to(device),
)
return replay_buffer
資料收集器 (Data collector)¶
與 PPO 和 DDPG 一樣，我們將使用資料收集器作為外層迴圈中的資料載入器。
我們選擇以下配置：我們將在一系列並行環境中同步並行執行，這些並行環境位於不同的收集器中，而這些收集器本身並行但非同步執行。
注意
此功能僅在 Python 多程序庫的“spawn”啟動方法中執行程式碼時可用。如果此教程直接作為指令碼執行（從而使用“fork”方法），我們將使用常規的 SyncDataCollector。
此配置的優點是我們可以在批次執行的計算量與我們希望非同步執行的計算量之間取得平衡。我們鼓勵讀者透過修改收集器數量（即傳遞給收集器的環境建構函式數量）以及每個收集器中並行執行的環境數量（由 num_workers 超引數控制）來實驗收集速度受到的影響。
收集器的裝置可以透過 device（通用）、policy_device、env_device 和 storing_device 引數完全引數化。storing_device 引數將修改正在收集的資料的位置：如果正在收集的批次大小相當可觀，我們可能希望將它們儲存在與計算發生位置不同的位置。對於我們的非同步資料收集器，不同的儲存裝置意味著我們收集的資料不會每次都位於同一裝置上，這是我們的訓練迴圈必須考慮到的。為簡單起見，我們將所有子收集器的裝置設定為相同的值。
def get_collector(
stats,
num_collectors,
actor_explore,
frames_per_batch,
total_frames,
device,
):
# We can't use nested child processes with mp_start_method="fork"
if is_fork:
cls = SyncDataCollector
env_arg = make_env(parallel=True, obs_norm_sd=stats, num_workers=num_workers)
else:
cls = MultiaSyncDataCollector
env_arg = [
make_env(parallel=True, obs_norm_sd=stats, num_workers=num_workers)
] * num_collectors
data_collector = cls(
env_arg,
policy=actor_explore,
frames_per_batch=frames_per_batch,
total_frames=total_frames,
# this is the default behavior: the collector runs in ``"random"`` (or explorative) mode
exploration_type=ExplorationType.RANDOM,
# We set the all the devices to be identical. Below is an example of
# heterogeneous devices
device=device,
storing_device=device,
split_trajs=False,
postproc=MultiStep(gamma=gamma, n_steps=5),
)
return data_collector
損失函式 (Loss function)¶
構建我們的損失函式很簡單：我們只需要向 DQNLoss 類提供模型和一些超引數。
目標引數¶
許多離策略 RL 演算法在使用“目標引數”來估計下一個狀態或狀態-動作對的值時會使用該概念。目標引數是模型引數的滯後副本。因為它們的預測與當前模型配置不匹配，它們透過對正在估計的值設定悲觀界限來幫助學習。這是一個強大的技巧（稱為“雙 Q-Learning”），在類似演算法中無處不在。
def get_loss_module(actor, gamma):
loss_module = DQNLoss(actor, delay_value=True)
loss_module.make_value_estimator(gamma=gamma)
target_updater = SoftUpdate(loss_module, eps=0.995)
return loss_module, target_updater
超引數¶
讓我們從超引數開始。以下設定在實踐中應該效果很好，並且演算法的效能應該不太對這些引數的微小變化敏感。
is_fork = multiprocessing.get_start_method() == "fork"
device = (
torch.device(0)
if torch.cuda.is_available() and not is_fork
else torch.device("cpu")
)
最佳化器¶
# the learning rate of the optimizer
lr = 2e-3
# weight decay
wd = 1e-5
# the beta parameters of Adam
betas = (0.9, 0.999)
# Optimization steps per batch collected (aka UPD or updates per data)
n_optim = 8
DQN 引數¶
gamma 衰減因子
gamma = 0.99
平滑目標網路更新衰減引數。這大致相當於具有硬目標網路更新的 1/tau 間隔
tau = 0.02
資料收集和回放緩衝區¶
注意
已註釋掉用於正確訓練的值。
在環境中收集的總幀數。在其他實現中，使用者定義了最大回合數。這在使用返回 N 幀批次的 Our data collectors 時更難做到，其中 N 是一個常數。但是，可以透過在收集了一定數量的回合後中斷訓練迴圈來輕鬆獲得相同數量的回合限制。
total_frames = 5_000 # 500000
用於初始化回放緩衝區的隨機幀。
init_random_frames = 100 # 1000
每個收集批次中的幀。
frames_per_batch = 32 # 128
每次最佳化步驟從回放緩衝區取樣
batch_size = 32 # 256
回放緩衝區的規模（以幀為單位）
buffer_size = min(total_frames, 100000)
每個資料收集器中並行執行的環境數量
num_workers = 2 # 8
num_collectors = 2 # 4
環境和探索¶
我們設定了 epsilon 貪婪探索中 epsilon 因子的初始值和最終值。由於我們的策略是確定性的，探索至關重要：沒有它，隨機性的唯一來源將是環境重置。
eps_greedy_val = 0.1
eps_greedy_val_env = 0.005
為了加速學習，我們將值網路的最後一層的偏差設定為預定義值（這不是強制性的）
init_bias = 2.0
注意
為了快速渲染本教程，total_frames 超引數被設定為非常低的值。要獲得合理的效能，請使用更大的值，例如 500000。
構建訓練器¶
TorchRL 的 Trainer 類建構函式接受以下僅關鍵字引數
collectorloss_moduleoptimizerlogger：Logger 可以是total_frames：此引數定義了訓練器的生命週期。frame_skip：當使用幀跳時，收集器必須知道它，以便準確計算收集的幀數等。讓訓練器知道此引數不是強制性的，但有助於在總幀數（預算）固定但幀跳可變的情況下進行更公平的比較。
stats = get_norm_stats()
test_env = make_env(parallel=False, obs_norm_sd=stats)
# Get model
actor, actor_explore = make_model(test_env)
loss_module, target_net_updater = get_loss_module(actor, gamma)
collector = get_collector(
stats=stats,
num_collectors=num_collectors,
actor_explore=actor_explore,
frames_per_batch=frames_per_batch,
total_frames=total_frames,
device=device,
)
optimizer = torch.optim.Adam(
loss_module.parameters(), lr=lr, weight_decay=wd, betas=betas
)
exp_name = f"dqn_exp_{uuid.uuid1()}"
tmpdir = tempfile.TemporaryDirectory()
logger = CSVLogger(exp_name=exp_name, log_dir=tmpdir.name)
warnings.warn(f"log dir: {logger.experiment.log_dir}")
我們可以控制標量多久記錄一次。這裡我們將其設定為一個較低的值，因為我們的訓練迴圈很短
log_interval = 500
trainer = Trainer(
collector=collector,
total_frames=total_frames,
frame_skip=1,
loss_module=loss_module,
optimizer=optimizer,
logger=logger,
optim_steps_per_batch=n_optim,
log_interval=log_interval,
)
註冊鉤子¶
可以透過兩種獨立的方式註冊鉤子
如果鉤子有它，
register()方法是首選。只需提供訓練器作為輸入，鉤子將以預設名稱在預設位置註冊。對於某些鉤子，註冊可能非常複雜：ReplayBufferTrainer需要 3 個鉤子（extend、sample和update_priority），這可能很麻煩。
buffer_hook = ReplayBufferTrainer(
get_replay_buffer(buffer_size, n_optim, batch_size=batch_size, device=device),
flatten_tensordicts=True,
)
buffer_hook.register(trainer)
weight_updater = UpdateWeights(collector, update_weights_interval=1)
weight_updater.register(trainer)
recorder = LogValidationReward(
record_interval=100, # log every 100 optimization steps
record_frames=1000, # maximum number of frames in the record
frame_skip=1,
policy_exploration=actor_explore,
environment=test_env,
exploration_type=ExplorationType.DETERMINISTIC,
log_keys=[("next", "reward")],
out_keys={("next", "reward"): "rewards"},
log_pbar=True,
)
recorder.register(trainer)
探索模組 epsilon 因子也會衰減
trainer.register_op("post_steps", actor_explore[1].step, frames=frames_per_batch)
任何可呼叫物件（包括
TrainerHookBase子類）都可以使用register_op()進行註冊。在這種情況下，必須顯式傳遞位置（）。此方法提供了對鉤子位置的更多控制，但也需要對訓練器機制有更多的瞭解。請參閱 訓練器文件 以獲取訓練器鉤子的詳細說明。
trainer.register_op("post_optim", target_net_updater.step)
我們也可以記錄訓練獎勵。請注意，對於 CartPole 來說，這興趣有限，因為獎勵總是 1。折扣獎勵總和的最佳化不是透過獲得更高的獎勵，而是透過讓 CartPole 存活更長時間來實現的。這將在進度條中顯示的 total_rewards 值中得到反映。
log_reward = LogScalar(log_pbar=True)
log_reward.register(trainer)
注意
如果需要，可以將多個最佳化器連結到訓練器。在這種情況下，每個最佳化器都將繫結到損失字典中的一個欄位。有關更多資訊，請參閱 OptimizerHook。
我們準備好訓練我們的演算法了！只需呼叫 trainer.train()，我們就會在日誌中看到我們的結果。
trainer.train()
我們現在可以快速檢視包含結果的 CSV 檔案。
def print_csv_files_in_folder(folder_path):
"""
Find all CSV files in a folder and prints the first 10 lines of each file.
Args:
folder_path (str): The relative path to the folder.
"""
csv_files = []
output_str = ""
for dirpath, _, filenames in os.walk(folder_path):
for file in filenames:
if file.endswith(".csv"):
csv_files.append(os.path.join(dirpath, file))
for csv_file in csv_files:
output_str += f"File: {csv_file}\n"
with open(csv_file) as f:
for i, line in enumerate(f):
if i == 10:
break
output_str += line.strip() + "\n"
output_str += "\n"
print(output_str)
print_csv_files_in_folder(logger.experiment.log_dir)
trainer.shutdown()
del trainer
結論和可能的改進¶
在本教程中，我們學習了
如何編寫一個訓練器，包括構建其元件並將它們註冊到訓練器中；
如何編寫一個 DQN 演算法，包括如何建立具有
QValueNetwork的策略來選擇具有最高值的動作；如何構建一個多程序資料收集器；
本教程的可能改進包括
也可以使用優先回放緩衝區。這將為具有最差值精度的樣本提供更高的優先順序。在文件的 回放緩衝區部分 中瞭解更多資訊。
分佈損失（有關更多資訊，請參閱
DistributionalDQNLoss）。更高階的探索技術，例如
NoisyLinear層等。
