注意

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

TorchRL 模块入门¶

注意

要在 notebook 中运行本教程，请在开头添加一个安装单元格，其中包含：

!pip install tensordict
!pip install torchrl

强化学习旨在创建能够有效应对特定任务的策略。策略可以采取多种形式，从将观测空间映射到动作空间的微分映射，到更临时的 F方法，例如对为每个可能动作计算的值列表取 argmax。策略可以是确定性的或随机的，并且可以包含复杂的元素，例如循环神经网络 (RNN) 或 Transformer。

容纳所有这些场景可能相当复杂。在本简洁的教程中，我们将深入探讨 TorchRL 在策略构建方面の核心功能。我们将主要关注两种常见场景下的随机策略和 Q 值策略：使用多层感知机 (MLP) 或卷积神经网络 (CNN) 作为骨干网络。

TensorDictModules¶

与环境与 `TensorDict` 实例交互的方式类似，用于表示策略和价值函数的模块也执行相同的操作。核心思想很简单：将标准的 `torch.nn.Module`（或任何其他函数）封装在一个类中，该类知道哪些条目需要被读取并传递给模块，然后使用指定的条目记录结果。为了说明这一点，我们将使用最简单的策略：从观测空间到动作空间的确定性映射。为了最大程度的通用性，我们将使用 Pendulum 环境的 `LazyLinear` 模块，该环境在上一教程中已实例化。

import torch

from tensordict.nn import TensorDictModule
from torchrl.envs import GymEnv

env = GymEnv("Pendulum-v1")
module = torch.nn.LazyLinear(out_features=env.action_spec.shape[-1])
policy = TensorDictModule(
    module,
    in_keys=["observation"],
    out_keys=["action"],
)

要执行我们的策略，仅此而已！懒惰模块的使用使我们无需获取观测空间的形状，因为模块将自动确定它。此策略现在已准备好在环境中运行。

rollout = env.rollout(max_steps=10, policy=policy)
print(rollout)

专用包装器¶

为了简化 `Actor`、# `ProbabilisticActor`、# `ActorValueOperator` 或 # `ActorCriticOperator` 的集成。例如，`Actor` 为 `in_keys` 和 `out_keys` 提供了默认值，使得与许多常见环境的集成变得简单。

from torchrl.modules import Actor

policy = Actor(module)
rollout = env.rollout(max_steps=10, policy=policy)
print(rollout)

可用专用 TensorDictModules 的列表可在 API 参考中找到。

网络¶

TorchRL 还提供了可以在不依赖 tensordict 功能的情况下使用的常规模块。您将遇到的最常见的两个网络是 `MLP` 和 `ConvNet` (CNN) 模块。我们可以用这些模块之一替换我们的策略模块。

from torchrl.modules import MLP

module = MLP(
    out_features=env.action_spec.shape[-1],
    num_cells=[32, 64],
    activation_class=torch.nn.Tanh,
)
policy = Actor(module)
rollout = env.rollout(max_steps=10, policy=policy)

TorchRL 还支持基于 RNN 的策略。由于这是一个更技术性的主题，因此在单独的教程中进行了介绍。

随机策略¶

像 PPO 这样的策略优化算法要求策略是随机的：与上面的示例不同，模块现在编码了一个从观测空间到编码动作空间上分布的参数空间的映射。TorchRL 通过将各种操作（例如从参数构建分布、从该分布采样以及检索对数概率）组合到一个类中来促进此类模块的设计。在这里，我们将构建一个依赖于具有三个组件的常规正态分布的 actor：

一个 `MLP` 主干，读取大小为 `[3]` 的观测值，并输出大小为 `[2]` 的单个张量；
一个 `NormalParamExtractor` 模块，它将此输出分成两个块：均值和标准差，大小均为 `[1]`；
一个 `ProbabilisticActor`，它将这些参数作为 `in_keys` 读取，创建具有这些参数的分布，并使用样本和对数概率填充我们的 tensordict。

from tensordict.nn.distributions import NormalParamExtractor
from torch.distributions import Normal
from torchrl.modules import ProbabilisticActor

backbone = MLP(in_features=3, out_features=2)
extractor = NormalParamExtractor()
module = torch.nn.Sequential(backbone, extractor)
td_module = TensorDictModule(module, in_keys=["observation"], out_keys=["loc", "scale"])
policy = ProbabilisticActor(
    td_module,
    in_keys=["loc", "scale"],
    out_keys=["action"],
    distribution_class=Normal,
    return_log_prob=True,
)

rollout = env.rollout(max_steps=10, policy=policy)
print(rollout)

关于此滚动的有几点需要注意：

由于我们在 actor 构建期间请求了它，因此给定当时分布的动作的对数概率也被写入。这对于 PPO 等算法是必需的。
分布的参数也作为 `loc` 和 `scale` 条目返回到输出 tensordict 中。

如果您的应用程序需要，您可以控制动作的采样，以使用分布的期望值或其他属性，而不是使用随机样本。这可以通过 `set_exploration_type()` 函数进行控制。

from torchrl.envs.utils import ExplorationType, set_exploration_type

with set_exploration_type(ExplorationType.DETERMINISTIC):
    # takes the mean as action
    rollout = env.rollout(max_steps=10, policy=policy)
with set_exploration_type(ExplorationType.RANDOM):
    # Samples actions according to the dist
    rollout = env.rollout(max_steps=10, policy=policy)

检查文档字符串中的 `default_interaction_type` 关键字参数以了解更多信息。

探索¶

像这样随机策略在一定程度上自然地权衡了探索和利用，但确定性策略则不能。幸运的是，TorchRL 也可以通过其探索模块来缓解这个问题。我们将以 `EGreedyModule` 探索模块为例（另请参阅 `AdditiveGaussianModule` 和 `OrnsteinUhlenbeckProcessModule`）。为了看到这个模块在起作用，让我们回到确定性策略。

from tensordict.nn import TensorDictSequential
from torchrl.modules import EGreedyModule

policy = Actor(MLP(3, 1, num_cells=[32, 64]))

我们的 ε-greedy 探索模块通常会使用一定的衰减帧数和 ε 参数的初始值进行自定义。ε = 1 的值意味着每个采取的动作都是随机的，而 ε=0 意味着根本没有探索。为了衰减（即减小）探索因子，需要调用 `step()`（有关示例，请参阅最后一个教程）。

exploration_module = EGreedyModule(
    spec=env.action_spec, annealing_num_steps=1000, eps_init=0.5
)

为了构建我们的探索性策略，我们只需要在 `TensorDictSequential` 模块（在 tensordict 领域中相当于 `Sequential`）内将确定性策略模块与探索模块串联起来。

exploration_policy = TensorDictSequential(policy, exploration_module)

with set_exploration_type(ExplorationType.DETERMINISTIC):
    # Turns off exploration
    rollout = env.rollout(max_steps=10, policy=exploration_policy)
with set_exploration_type(ExplorationType.RANDOM):
    # Turns on exploration
    rollout = env.rollout(max_steps=10, policy=exploration_policy)

由于它必须能够从动作空间中采样随机动作，因此 `EGreedyModule` 必须配备环境的 `action_space` 才能知道使用什么策略来随机采样动作。

Q 值 actor¶

在某些情况下，策略不是一个独立的模块，而是构建在另一个模块之上。Q 值 actor 就是这种情况。简而言之，这些 actor 需要动作值的估计（大多数时候是离散的），并且会贪婪地选择具有最高值的动作。在某些情况下（有限离散动作空间和有限离散状态空间），可以只存储一个状态-动作对的 2D 表格，并选择值最高的动作。DQN 带来的创新是通过使用神经网络来编码 `Q(s, a)` 值映射，将此扩展到连续状态空间。为了更清楚地理解，让我们考虑另一个具有离散动作空间的环境。

env = GymEnv("CartPole-v1")
print(env.action_spec)

我们构建了一个价值网络，当它从环境中读取状态时，为每个动作产生一个值。

num_actions = 2
value_net = TensorDictModule(
    MLP(out_features=num_actions, num_cells=[32, 32]),
    in_keys=["observation"],
    out_keys=["action_value"],
)

我们可以通过在我们的价值网络之后添加一个 `QValueModule` 来轻松构建我们的 Q 值 actor。

from torchrl.modules import QValueModule

policy = TensorDictSequential(
    value_net,  # writes action values in our tensordict
    QValueModule(spec=env.action_spec),  # Reads the "action_value" entry by default
)

让我们来看看！我们让策略运行几个步骤，并查看输出。我们应该在滚动中找到 `"action_value"` 以及 `"chosen_action_value"` 条目。

rollout = env.rollout(max_steps=3, policy=policy)
print(rollout)

由于它依赖于 `argmax` 运算符，因此此策略是确定性的。在数据收集期间，我们需要探索环境。为此，我们再次使用 `EGreedyModule`。

policy_explore = TensorDictSequential(policy, EGreedyModule(env.action_spec))

with set_exploration_type(ExplorationType.RANDOM):
    rollout_explore = env.rollout(max_steps=3, policy=policy_explore)

这就是我们关于使用 TorchRL 构建策略的简短教程！

您还可以使用该库进行许多其他操作。一个好的起点是查看模块 API 参考。

后续步骤

查看当动作是复合的时如何使用 `CompositeDistribution`（例如，env 需要离散和连续动作）；
看看如何在策略中使用 RNN（教程）；
将此与 Decision Transformers 示例中的 Transformer 使用进行比较（参见 GitHub 上的 `example` 目录）。

由 Sphinx-Gallery 生成的画廊