注意

跳转到末尾下载完整的示例代码。

使用 TorchRL 训练强化学习 (PPO) 教程¶

作者：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 通常被认为是用于在线、on-policy 强化算法的快速高效方法。TorchRL 提供了一个损失模块，可以为您完成所有工作，因此您可以依赖此实现，专注于解决问题，而不是每次想要训练策略时都重新发明轮子。

为完整起见，这里是对损失计算的简要概述，尽管这由我们的 ClipPPOLoss 模块处理—算法如下： 1. 我们将通过在环境中执行策略给定的步数来采样一个数据批。 2. 然后，我们将使用裁剪版的 REINFORCE 损失对该批数据进行随机子采样，并执行给定次数的优化步骤。 3. 裁剪将对我们的损失设置一个悲观下界：相比于较高的回报估计，较低的回报估计将受到青睐。损失的确切公式为

\[L(s,a,\theta_k,\theta) = \min\left( \frac{\pi_{\theta}(a|s)}{\pi_{\theta_k}(a|s)} A^{\pi_{\theta_k}}(s,a), \;\; g(\epsilon, A^{\pi_{\theta_k}}(s,a)) \right),\]

该损失有两个组成部分：在最小运算符的第一部分，我们简单地计算 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 参数¶

在每次数据收集（或批次收集）时，我们将在一定数量的 *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 控制套件等。作为通用库，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)¶

我们将向我们的环境添加一些变换，以准备策略的数据。在 Gym 中，这通常通过包装器实现。TorchRL 采用不同的方法，更类似于其他 PyTorch 领域库，通过使用变换。要向环境添加变换，只需将其包装在 TransformedEnv 实例中，并附加变换序列。转换后的环境将继承被包装环境的设备和元数据，并根据其包含的变换序列转换这些内容。

标准化¶

首先要编码的是标准化变换。经验法则规定，最好使数据大致匹配单位高斯分布：为了实现这一点，我们将执行一定数量的随机步骤，并计算这些观察的统计数据。

我们将添加另外两个变换：DoubleToFloat 变换会将双精度条目转换为单精度数字，以便策略读取。StepCounter 变换将用于计算环境终止前的步数。我们将使用此度量作为性能的补充度量。

正如我们稍后将看到的，TorchRL 的许多类都依赖 TensorDict 进行通信。您可以将其视为具有一些额外张量功能的 Python 字典。实际上，这意味着我们将处理的许多模块都需要被告知读取哪个键（in_keys）以及写入哪个键（out_keys）在它们将收到的 tensordict 中。通常，如果省略 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() 函数运行一个小的回滚，并将其输出与环境规范进行比较。如果没有引发错误，我们可以确信规范已正确定义。

check_env_specs(env)

为了好玩，让我们看看一个简单的随机滚动是什么样的。您可以调用 env.rollout(n_steps) 并大致了解环境输入和输出的样子。动作将自动从动作规范域中抽取，因此您不必担心设计一个随机采样器。

通常，在每一步，RL 环境都会接收一个动作作为输入，并输出一个观察、一个奖励和一个完成状态。观察可能是复合的，这意味着它可能由多个张量组成。这对于 TorchRL 来说不是问题，因为整个观察集会自动打包到输出的 TensorDict 中。在执行给定步数的滚动（例如，一系列环境步骤和随机动作生成）后，我们将检索一个 TensorDict 实例，其形状与此轨迹长度匹配

rollout = env.rollout(3)
print("rollout of three steps:", rollout)
print("Shape of the rollout TensorDict:", rollout.batch_size)

我们的滚动数据形状为 torch.Size([3])，与我们运行的步数匹配。"next" 条目指向当前步骤之后的数据。在大多数情况下，时间 t 的 "next" 数据与 t+1 的数据匹配，但如果使用了某些特定变换（例如，多步），则可能不匹配。

策略 (Policy)¶

PPO 使用随机策略来处理探索。这意味着我们的神经网络将不得不输出一个分布的参数，而不是对应于所采取动作的单个值。

由于数据是连续的，我们使用 Tanh-Normal 分布来尊重动作空间边界。TorchRL 提供了这种分布，我们唯一需要关心的是构建一个输出正确数量参数的神经网络，以便策略可以使用（一个位置，或均值，和一个缩放因子）

\[f_{\theta}(\text{observation}) = \mu_{\theta}(\text{observation}), \sigma^{+}_{\theta}(\text{observation})\]

这里带来的唯一额外困难是将我们的输出分成两部分，并将第二部分映射到一个严格正的空间。

我们分三步设计策略

定义一个神经网络 D_obs -> 2 * D_action。事实上，我们的 loc（均值）和 scale（标准差）都具有 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_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()))

数据收集器 (Data collector)¶

TorchRL 提供了一组数据收集器类。简而言之，这些类执行三个操作：重置环境，根据最新观察计算动作，在环境中执行一步，然后重复最后两个步骤，直到环境发出停止信号（或达到完成状态）。

它们允许您控制每次迭代收集多少帧（通过 frames_per_batch 参数），何时重置环境（通过 max_frames_per_traj 参数），策略应在哪个 device 上执行，等等。它们还旨在与批量和多进程环境高效地协同工作。

最简单的数据收集器是 SyncDataCollector：它是一个迭代器，您可以使用它来获取给定长度的数据批，并在收集完总帧数（total_frames）后停止。其他数据收集器（MultiSyncDataCollector 和 MultiaSyncDataCollector）将在多进程工作者集合上以同步和异步方式执行相同的操作。

与之前的策略和环境一样，数据收集器将返回 TensorDict 实例，其总元素数量将匹配 frames_per_batch。使用 TensorDict 将数据传递给训练循环允许您编写 100% 忽略滚动内容实际特异性的数据加载管道。

collector = SyncDataCollector(
    env,
    policy_module,
    frames_per_batch=frames_per_batch,
    total_frames=total_frames,
    split_trajs=False,
    device=device,
)

回放缓冲区 (Replay buffer)¶

回放缓冲区是离策略 RL 算法的常见构建模块。在策略环境中，每当收集一批数据时，回放缓冲区就会被重新填充，并且其数据会在一定数量的 epoch 中被重复消耗。

TorchRL 的回放缓冲区是使用通用容器 ReplayBuffer 构建的，它接受缓冲区的组件作为参数：存储、写入器、采样器以及可能的变换。只有存储（指示回放缓冲区容量）是必需的。我们还指定了一个无重复采样器，以避免在一个 epoch 中多次采样同一项。将回放缓冲区用于 PPO 不是必需的，我们可以简单地从收集的数据批中采样子批次，但使用这些类可以轻松地以可重现的方式构建内部训练循环。

replay_buffer = ReplayBuffer(
    storage=LazyTensorStorage(max_size=frames_per_batch),
    sampler=SamplerWithoutReplacement(),
)

损失函数 (Loss function)¶

PPO 损失可以方便地从 TorchRL 中直接导入，使用 ClipPPOLoss 类。这是使用 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
)

训练循环¶

现在我们有了编写训练循环所需的所有组件。步骤包括

收集数据
- 计算优势
  - 循环遍历收集的数据以计算损失值
  - 反向传播
  - 优化
  - 重复
- 重复
重复

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()

结果 (Results)¶

在达到 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。

由 Sphinx-Gallery 生成的画廊