强化学习 (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 通常被认为是在线、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 损失的带重要性权重的版本（例如，我们已经修正了当前策略配置滞后于数据收集时使用的配置的事实）。此最小运算符的第二部分是一个类似的损失，其中我们将比例因子在超过或低于给定阈值对时进行了剪辑。

此损失确保了无论优势是正数还是负数，都会抑制那些会导致与先前配置发生重大变化的策略更新。

本教程结构如下

1. 首先，我们将定义一组用于训练的超参数。

2. 接下来，我们将专注于使用 TorchRL 的包装器和转换器来创建我们的环境或模拟器。

3. 接下来，我们将设计策略网络和价值模型，这对于损失函数是必不可少的。这些模块将用于配置我们的损失模块。

4. 接下来，我们将创建回放缓冲区和数据加载器。

5. 最后，我们将运行我们的训练循环并分析结果。

在本教程中，我们将使用 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 参数

在每次数据收集（或批次收集）时，我们将在一定数量的 *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

定义环境

在 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})\]

这里唯一增加的难度是分割我们的输出成两部分，并将第二部分映射到一个严格正的空间。

我们分三步设计策略

1. 定义一个神经网络 D_obs -> 2 * D_action。事实上，我们的 loc（mu）和 scale（sigma）都具有 D_action 的维度。

2. 附加一个 NormalParamExtractor 来提取位置和尺度（例如，将输入分成相等的两部分，并将尺度参数映射到一个正值）。

3. 创建一个概率性的 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()

结论和后续步骤

在本教程中，我们学习了

1. 如何使用 torchrl 创建和自定义环境；

2. 如何编写模型和损失函数；

3. 如何设置典型的训练循环。

如果您想进一步尝试此教程，可以进行以下修改

• 从效率的角度来看，我们可以并行运行多个模拟来加速数据收集。有关更多信息，请查看 ParallelEnv。

• 从日志记录的角度来看，我们可以将 torchrl.record.VideoRecorder 转换器附加到环境，在请求渲染后，以获得倒立摆动作的视觉渲染。有关更多信息，请查看 torchrl.record。