强化学习（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

Proximal Policy Optimization (PPO) 是一种策略梯度算法，其中收集一批数据并直接用于训练策略以最大化在某些接近性约束下的预期回报。你可以将其视为 REINFORCE 的复杂版本，REINFORCE 是基础的策略优化算法。更多信息，请参阅Proximal Policy Optimization Algorithms 论文。

PPO 通常被视为一种快速且高效的在线、策略内强化学习算法。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 损失）。 该最小值算子的第二部分则是一种类似的损失，其中当比率超出或低于给定的一对阈值时，我们对其进行了裁剪。

该损失函数确保：无论优势值为正或负，都会抑制那些可能导致策略从先前配置发生显著变化的更新。

本教程结构如下：

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

2. 接下来，我们将重点利用 TorchRL 的封装器（wrappers）和变换（transforms）来构建我们的环境（即模拟器）。

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的目标是提供一个可互换的接口，以便您可以轻松地切换一个环境为另一个环境。例如，创建一个包装的gym环境只需几行字符：

base_env = GymEnv("InvertedDoublePendulum-v4", device=device)

在这段代码中需要注意几件事：首先，我们通过调用GymEnv包装器创建了环境。如果传递了额外的关键字参数，它们将会被传送到gym.make方法中，从而覆盖最常见的环境构建命令。 或者，也可以直接使用gym.make(env_name, **kwargs)创建一个gym环境，并将其封装在GymWrapper类中。

Also the device参数：对于gym，这仅控制输入动作和观察状态将被存储的设备，但执行始终在CPU上进行。这样做的原因是gym不支持设备上的执行，除非另有指定。对于其他库，我们对执行设备有控制权，并尽可能在存储和执行后端方面保持一致。

变换

我们将为环境添加一些转换以准备数据供策略使用。在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)

The ObservationNorm 变换现在已填充了位置和比例，这些将用于归一化数据。

让我们对总结统计量的形状进行一个小的合理性检查：

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

normalization constant shape: torch.Size([11])

一个环境不仅由其模拟器和转换器定义，还由一系列描述在执行过程中可以预期到的元数据定义。 为了提高效率，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)

observation_spec: Composite(
    observation: UnboundedContinuous(
        shape=torch.Size([11]),
        space=ContinuousBox(
            low=Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, contiguous=True),
            high=Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, contiguous=True)),
        device=cpu,
        dtype=torch.float32,
        domain=continuous),
    step_count: BoundedDiscrete(
        shape=torch.Size([1]),
        space=ContinuousBox(
            low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True),
            high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True)),
        device=cpu,
        dtype=torch.int64,
        domain=discrete),
    device=cpu,
    shape=torch.Size([]))
reward_spec: UnboundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)
input_spec: Composite(
    full_state_spec: Composite(
        step_count: BoundedDiscrete(
            shape=torch.Size([1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True),
                high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, contiguous=True)),
            device=cpu,
            dtype=torch.int64,
            domain=discrete),
        device=cpu,
        shape=torch.Size([])),
    full_action_spec: Composite(
        action: BoundedContinuous(
            shape=torch.Size([1]),
            space=ContinuousBox(
                low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
                high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
            device=cpu,
            dtype=torch.float32,
            domain=continuous),
        device=cpu,
        shape=torch.Size([])),
    device=cpu,
    shape=torch.Size([]))
action_spec (as defined by input_spec): BoundedContinuous(
    shape=torch.Size([1]),
    space=ContinuousBox(
        low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
        high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
    device=cpu,
    dtype=torch.float32,
    domain=continuous)

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

rollout of three steps: TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 11]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                step_count: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.int64, is_shared=False),
                terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([3]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 11]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([3]),
    device=cpu,
    is_shared=False)
Shape of the rollout TensorDict: torch.Size([3])

我们的发布数据形状为 torch.Size([3])，这与我们运行的步数相符。"next" 进入点指向当前步骤之后的数据。 在大多数情况下，时间 t 的 "next" 数据与 t+1 处的数据匹配，但如果使用了一些特定的变换（例如多步变换），这可能不成立。

政策

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. Append a 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()))

Running policy: TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        loc: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        observation: Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
Running value: TensorDict(
    fields={
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([11]), device=cpu, dtype=torch.float32, is_shared=False),
        state_value: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        step_count: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

数据收集器

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

它们允许你在每次迭代中控制收集多少帧（通过frames_per_batch参数）， 何时重置环境（通过max_frames_per_traj参数）， 在哪个device上执行策略等。它们还设计为能够高效地与批处理和多进程环境配合使用。

最简单的数据收集器是 SyncDataCollector： 它是一个迭代器，你可以用它来获取给定长度的数据批次，并且 一旦总共收集了指定数量的帧 (total_frames) 就会停止。 其他数据收集器 (MultiSyncDataCollector 和 MultiaSyncDataCollector) 将在一组多进程工作进程中以同步和异步方式执行 相同的操作。

在之前的策略和环境中，数据收集器将返回 TensorDict 个实例，并且元素总数将匹配 frames_per_batch。使用 TensorDict 将数据传递给训练循环可以让你编写完全不知情于实际展开内容具体细节的数据加载管道。

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

回放缓冲区

经验回放缓冲区（Replay buffers）是无策略（off-policy）强化学习算法中常见的组成部分。 在有策略（on-policy）场景中，每次收集一批数据后都会重新填充经验回放缓冲区，且其数据会在一定数量的训练轮次（epochs）内被反复使用。

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

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

损失函数

为方便起见，PPO 损失函数可直接从 TorchRL 导入，使用的是 ClipPPOLoss 类。这是使用 PPO 的最简单方式： 它隐藏了 PPO 的数学运算以及与之相关的控制流程。

PPO 需要计算一些“优势估计”。简而言之，优势是一个值，反映了在处理偏差/方差权衡时对回报值的期望。 为了计算优势，只需要（1）构建优势模块，该模块利用我们的价值运算符，并且（2）在每个周期之前将每批数据通过它。 GAE 模块将更新输入 tensordict 中的新 "advantage" 和 "value_target" 条目。 "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()

  0%|          | 0/10000 [00:00<?, ?it/s]
 10%|█         | 1000/10000 [00:02<00:25, 352.51it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.0788 (init= 9.0788), step count (max): 11, lr policy:  0.0003:  10%|█         | 1000/10000 [00:02<00:25, 352.51it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.0788 (init= 9.0788), step count (max): 11, lr policy:  0.0003:  20%|██        | 2000/10000 [00:05<00:22, 351.26it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.0946 (init= 9.0788), step count (max): 17, lr policy:  0.0003:  20%|██        | 2000/10000 [00:05<00:22, 351.26it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.0946 (init= 9.0788), step count (max): 17, lr policy:  0.0003:  30%|███       | 3000/10000 [00:08<00:19, 353.23it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.1386 (init= 9.0788), step count (max): 13, lr policy:  0.0003:  30%|███       | 3000/10000 [00:08<00:19, 353.23it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.1386 (init= 9.0788), step count (max): 13, lr policy:  0.0003:  40%|████      | 4000/10000 [00:11<00:17, 352.13it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.1905 (init= 9.0788), step count (max): 19, lr policy:  0.0002:  40%|████      | 4000/10000 [00:11<00:17, 352.13it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.1905 (init= 9.0788), step count (max): 19, lr policy:  0.0002:  50%|█████     | 5000/10000 [00:14<00:14, 352.48it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2068 (init= 9.0788), step count (max): 27, lr policy:  0.0002:  50%|█████     | 5000/10000 [00:14<00:14, 352.48it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2068 (init= 9.0788), step count (max): 27, lr policy:  0.0002:  60%|██████    | 6000/10000 [00:16<00:11, 353.99it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2200 (init= 9.0788), step count (max): 28, lr policy:  0.0001:  60%|██████    | 6000/10000 [00:16<00:11, 353.99it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2200 (init= 9.0788), step count (max): 28, lr policy:  0.0001:  70%|███████   | 7000/10000 [00:19<00:08, 355.93it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2301 (init= 9.0788), step count (max): 26, lr policy:  0.0001:  70%|███████   | 7000/10000 [00:19<00:08, 355.93it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2301 (init= 9.0788), step count (max): 26, lr policy:  0.0001:  80%|████████  | 8000/10000 [00:22<00:05, 357.78it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2338 (init= 9.0788), step count (max): 34, lr policy:  0.0001:  80%|████████  | 8000/10000 [00:22<00:05, 357.78it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2338 (init= 9.0788), step count (max): 34, lr policy:  0.0001:  90%|█████████ | 9000/10000 [00:25<00:02, 358.86it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2383 (init= 9.0788), step count (max): 26, lr policy:  0.0000:  90%|█████████ | 9000/10000 [00:25<00:02, 358.86it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2383 (init= 9.0788), step count (max): 26, lr policy:  0.0000: 100%|██████████| 10000/10000 [00:28<00:00, 359.95it/s]
eval cumulative reward:  82.1623 (init:  82.1623), eval step-count: 8, average reward= 9.2363 (init= 9.0788), step count (max): 33, lr policy:  0.0000: 100%|██████████| 10000/10000 [00:28<00:00, 359.95it/s]

结果

在达到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 以了解更多。

估计内存使用量： 318 MB