注意
转到末尾 下载完整示例代码。
强化学习 (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了解更多信息。
