备注

点击此处下载完整示例代码

使用TorchRL的强化学习（PPO）教程¶

Created On: Mar 15, 2023 | Last Updated: Mar 20, 2025 | Last Verified: Nov 05, 2024

本教程演示了如何使用PyTorch和:py:mod:`torchrl`训练一个参数化策略网络解决来自`OpenAI-Gym/Farama-Gymnasium控制库 <https://github.com/Farama-Foundation/Gymnasium>`__的倒立摆任务。

倒立摆¶

关键学习点：

如何在TorchRL中创建环境、转换其输出并从该环境中收集数据；
如何使用:class:`~tensordict.TensorDict`让您的类相互通信；
使用TorchRL构建训练循环的基础知识：
- 如何为策略梯度方法计算优势信号；
- 如何使用概率神经网络创建随机策略；
- 如何创建动态重放缓冲区并从中采样而无重复。

我们将涵盖TorchRL的六个关键组件：

如果您在 Google Colab 中运行此程序，请确保安装以下依赖项：

!pip3 install torchrl
!pip3 install gym[mujoco]
!pip3 install tqdm

近端策略优化（PPO）是一种策略梯度算法，先收集一批数据然后直接用于训练策略以在一定的近端性约束下最大化期望回报。您可以将其视为`REINFORCE <https://link.springer.com/content/pdf/10.1007/BF00992696.pdf>`_（基础策略优化算法）的一个复杂版本。更多信息请参阅`近端策略优化算法 <https://arxiv.org/abs/1707.06347>`_论文。

PPO通常被认为是一种快速高效的在线、基于策略的强化算法。TorchRL提供了一个无需您再做任何工作即可运行的损失模块，因此您可以依靠此实现专注于解决问题，而无需每次训练策略时重新发明轮子。

为了完整性，这里是对损失计算的简要概述，尽管这些已由我们的:class:`~torchrl.objectives.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损失（例如，我们修正了当前策略配置与用于收集数据时策略配置之间的迟滞差异）。最小函数的第二部分是一个相似的损失，但我们在比给定阈值高或低时对比率进行了修剪。

此损失可确保无论优势为正还是负，都会阻止产生重大偏移于之前配置的策略更新。

本教程结构如下：

首先，我们将定义训练中使用的一组超参数。
接下来，我们将专注于使用TorchRL的封装器和转换器来创建我们的环境或模拟器。
然后，我们将设计策略网络和价值模型，这对于损失函数来说是必不可少的。这些模块将用于配置我们的损失模块。
接下来，我们将创建重放缓冲区和数据加载器。
最后，我们将运行训练循环并分析结果。

在本教程中，我们将使用:mod:`tensordict`库。:class:`~tensordict.TensorDict`是TorchRL的通用语言：它帮助我们抽象模块的读写内容，减少对具体数据描述的关注，并更多关注算法本身。

import warnings
warnings.filterwarnings("ignore")
from torch import multiprocessing


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 = 50_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

定义一个环境¶

在强化学习中，称为*环境*的通常指模拟器或控制系统。各种库为强化学习提供了模拟环境，包括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不支持设备上的执行，除非另行指定。对于其他库，我们可以控制执行设备，并尽可能在存储和执行后端保持一致性。

转换¶

我们将给环境附加一些转换来为策略准备数据。在Gym中，这通常通过封装器实现。TorchRL采用了一种不同的方法，更类似于其他PyTorch领域库，通过使用转换器。为了向环境添加转换，应简单地将其封装在:class:`~torchrl.envs.transforms.TransformedEnv`实例中，并将转换序列附加到其上。转换后的环境将继承封装环境的设备和元数据，并根据其包含的转换序列转换这些属性。

归一化¶

首先编码的是归一化转换。通常来说，建议数据与单位高斯分布大致匹配：为此，我们将在环境中运行若干随机步骤并计算这些观察值的统计摘要。

我们将附加另外两个转换：:class:`~torchrl.envs.transforms.DoubleToFloat`转换将双精度条目转换为单精度数，以便于策略读取。而:class:`~torchrl.envs.transforms.StepCounter`转换将用于计算环境终止之前的步骤数量。我们将使用此测量作为性能的补充测度。

正如我们稍后将看到的，许多TorchRL类依赖于:class:~tensordict.TensorDict`进行通信。可以将其视为一个具有额外张量功能的Python字典。实际上，这意味着我们使用的许多模块需要指定从`tensordict`中读取的键(``in_keys`)以及写入的键(out_keys)。通常，如果省略``out_keys``，则假定更新``in_keys``条目。对于我们的转换层，我们只感兴趣一个称为``”observation”``的条目，并指示我们的转换层仅修改该条目：

env = TransformedEnv(
    base_env,
    Compose(
        # normalize observations
        ObservationNorm(in_keys=["observation"]),
        DoubleToFloat(),
        StepCounter(),
    ),
)

正如您可能注意到的，我们创建了一个归一化层，但是尚未设置其归一化参数。为此，:class:`~torchrl.envs.transforms.ObservationNorm`可以自动收集我们环境的汇总统计数据：

env.transform[0].init_stats(num_iter=1000, reduce_dim=0, cat_dim=0)

:class:`~torchrl.envs.transforms.ObservationNorm`转换现在已填充用于归一化数据的位置和比例。

让我们做一个小的检查，验证我们的汇总统计的形状：

print("normalization constant shape:", env.transform[0].loc.shape)

一个环境不仅由其模拟器和转换层定义，还由描述其执行过程中可以预期内容的一系列元数据定义。为了提高效率，TorchRL对环境规格有一定的严格要求，但可以很容易地检查您的环境规格是否合适。在我们的示例中，继承自:class:`~torchrl.envs.libs.gym.GymWrapper`和:class:`~torchrl.envs.libs.gym.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)

:func:`check_env_specs`函数通过小规模的试运行并将其输出与环境规格进行比较。如果没有出现错误，我们可以确信这些规格定义是正确的：

check_env_specs(env)

为了趣味，让我们看看一个简单的随机试运行会是什么样子。可以调用`env.rollout(n_steps)`并查看环境输入和输出的概览。动作将自动从动作规格域中抽取，因此不需要设计随机采样器。

通常，在每一步中，RL环境接收一个动作作为输入，并输出一个观察值、奖励以及完成状态。观察值可能是复合的，这意味着它可能由多个张量组成。这对于TorchRL来说不是问题，因为整个观察值集会自动打包到输出:class:`~tensordict.TensorDict`中。在执行试运行（例如，环境步和随机动作生成的序列）后，我们会获取一个:class:`~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``的数据一致，但如果我们使用某些特定的转换（例如多步），可能不会如此。

策略¶

PPO利用一种随机策略来进行探索。这意味着我们的神经网络需要输出一个分布的参数，而不是一个单一的值来表示所采取的动作。

由于数据是连续性，我们使用Tanh-正态分布来遵守动作空间的边界。TorchRL提供了这样的分布，我们需要关注的仅仅是构建能够输出正确数量参数的神经网络，以供策略使用（即位置或均值，以及一个比例）：

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

此处唯一的额外难点是将输出拆分为两部分，并将第二部分映射到一个严格正值的空间。

我们分三步设计策略：

定义一个神经网络 D_obs -> 2 * D_action。实际上，我们的``loc``（均值）和``scale``（标准差）都具有维度``D_action``。
附加一个:class:`~tensordict.nn.distributions.NormalParamExtractor`用于提取位置和比例（例如，将输入拆分为两部分，并对比例参数应用一个正值转换）。
创建一个概率性的:class:~tensordict.nn.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``包装在:class:~tensordict.nn.TensorDictModule`中。这个类只需要读取给定的``in_keys``并将输出写入到注册的``out_keys`。

policy_module = TensorDictModule(
    actor_net, in_keys=["observation"], out_keys=["loc", "scale"]
)

现在我们需要基于正态分布的位置和比例创建一个分布。为此，我们指示:class:~torchrl.modules.tensordict_module.ProbabilisticActor`类构建一个:class:`~torchrl.modules.TanhNormal，并提供该分布的最小值和最大值，这些值可以从环境规格中获取。

in_keys``的名称（以及上面:class:`~tensordict.nn.TensorDictModule`的``out_keys``名称）不能随意设置，因为:class:`~torchrl.modules.TanhNormal`分布构造函数将需要``loc``和``scale``关键词参数。不过，:class:`~torchrl.modules.tensordict_module.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"],
)

让我们尝试一下我们的策略和价值模块。正如我们之前所说，使用:class:`~tensordict.nn.TensorDictModule`使这两个模块可以直接读取环境的输出进行操作，因为它们知道要读取什么信息以及将写入何处：

print("Running policy:", policy_module(env.reset()))
print("Running value:", value_module(env.reset()))

数据收集器¶

TorchRL提供了一组`DataCollector类 <https://pytorch.org/rl/reference/collectors.html>`__。简而言之，这些类执行三个操作：重置环境，根据最新观察值计算动作，在环境中执行步骤，然后重复后两个步骤直到环境发送停止信号（或达到完成状态）。

这些类允许您控制每轮收集多少帧（通过``frames_per_batch``参数）、何时重置环境（通过``max_frames_per_traj``参数）、策略应该在哪个``device``上执行等。它们还旨在与批处理和多进程环境高效地协作。

最简单的数据收集器是:class:~torchrl.collectors.collectors.SyncDataCollector：它是一个迭代器，可以用于获取给定长度的数据批，并在收集总帧数（total_frames）后停止。其他数据收集器（MultiaSyncDataCollector）则在一组多进程工作者中以同步和异步方式执行相同操作。

与之前的策略和环境一样，数据收集器将返回与``frames_per_batch``匹配总元素数的:class:`~tensordict.TensorDict`实例。使用:class:`~tensordict.TensorDict`将数据传递到训练循环中允许您编写完全独立于试运行内容具体方面的数据加载管道。

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

重放缓冲区¶

重放缓冲区是离线策略RL算法中常见的构建块。在在线策略上下文中，重放缓冲区会在每次收集数据批时重新填充，并在一定次数的epoch中重复使用其数据。

TorchRL的重放缓冲区是使用:class:`~torchrl.data.ReplayBuffer`作为通用容器构建的，该容器接受缓冲区的组件：存储器、写入器、采样器以及可能的转换器。只有存储器（表示重放缓冲区容量）是必需的。我们还指定一个无重复采样器，以避免在一个epoch中对同一项进行多次采样。对于PPO使用重放缓冲区不是必须的，也可以直接从收集的数据中采样子批，但使用这些类使我们能够以更具复现性的方式构建内部训练循环。

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

损失函数¶

可以通过使用:class:`~torchrl.objectives.ClipPPOLoss`类直接从TorchRL导入PPO损失以方便使用。这是使用PPO的最简单方式：它隐藏了PPO的数学运算和相关的控制流程。

PPO需要计算一些”优势估计”。简而言之，优势是一种反映回报值期望的值，同时处理偏差与方差之间的权衡。为了计算优势，只需（1）构造使用我们的价值运算符的优势模块，然后（2）在每个epoch之前将每批数据传递给它。GAE模块会用新的``”advantage”和”value_target”条目更新输入的``tensordict。``”value_target”``是一个无梯度张量，表示价值网络在输入观察条件下应该预测的经验价值。这两者都将被:class:`~torchrl.objectives.ClipPPOLoss`用于返回策略和价值损失。

advantage_module = GAE(
    gamma=gamma, lmbda=lmbda, value_network=value_module, average_gae=True, device=device,
)

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 了解更多信息。

**脚本的总运行时间：**（0分钟0.000秒）

通过Sphinx-Gallery生成的图集