注意

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

TorchRL 模块入门¶

注意

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

!pip install tensordict
!pip install torchrl

强化学习旨在创建能够有效解决特定任务的策略。策略可以采用多种形式，从将观察空间转换为动作空间的可微映射，到更灵活的方法，例如对为每个可能动作计算的值列表进行 argmax。策略可以是确定性的或随机的，并且可以包含复杂的元素，例如循环神经网络 (RNN) 或变换器。

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

TensorDict 模块¶

与环境如何与 TensorDict 的实例交互类似，用于表示策略和价值函数的模块也这样做。核心思想很简单：将标准 Module（或任何其他函数）封装在一个知道需要读取哪些条目并将其传递给模块的类中，然后使用分配的条目记录结果。为了说明这一点，我们将使用最简单的策略：从观察空间到动作空间的确定性映射。为了最大限度地通用，我们将使用一个 LazyLinear 模块，以及我们在上一个教程中实例化的 Pendulum 环境。

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)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

专用包装器¶

为了简化 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)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

在 API 参考中提供了可用的专用 TensorDict 模块列表。

网络¶

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 通过在一个类中对各种操作进行分组来简化此类模块的设计，这些操作例如从参数构建分布、从该分布采样以及检索对数概率。在这里，我们将构建一个依靠正态分布的演员，使用三个组件

一个 MLP 骨干网络，读取大小为 [3] 的观察结果并输出一个大小为 [2] 的单个张量；
一个 NormalParamExtractor 模块，它将此输出分成两个部分，一个大小为 [1] 的均值和一个大小为 [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)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        loc: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([10]),
            device=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([10, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([10]),
    device=None,
    is_shared=False)

关于此回滚需要注意几点

由于我们在构建演员时请求了它，因此也会写入给定该时间分布的动作的对数概率。这对于像 PPO 这样的算法是必要的。
分布的参数也会在输出 tensordict 中返回，位于 "loc" 和 "scale" 条目下。

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

from torchrl.envs.utils import ExplorationType, set_exploration_type

with set_exploration_type(ExplorationType.MEAN):
    # 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]))

我们的\(\epsilon\)贪婪探索模块通常会根据退火帧数和\(\epsilon\)参数的初始值进行定制。\(\epsilon = 1\)的值表示每个采取的动作都是随机的，而\(\epsilon=0\)表示根本没有探索。为了退火（即减少）探索因子，需要调用step()（有关示例，请参见最后一个教程）。

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

为了构建我们的探索性策略，我们只需要将确定性策略模块与探索模块连接在一个TensorDictSequential模块中（这类似于张量字典领域中的Sequential）。

exploration_policy = TensorDictSequential(policy, exploration_module)

with set_exploration_type(ExplorationType.MEAN):
    # 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-Value 演员¶

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

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

OneHotDiscreteTensorSpec(
    shape=torch.Size([2]),
    space=DiscreteBox(n=2),
    device=cpu,
    dtype=torch.int64,
    domain=discrete)

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

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

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

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)

TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 2]), device=cpu, dtype=torch.int64, is_shared=False),
        action_value: Tensor(shape=torch.Size([3, 2]), device=cpu, dtype=torch.float32, is_shared=False),
        chosen_action_value: 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, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, 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=None,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, 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=None,
    is_shared=False)

由于它依赖于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，当动作是复合的时（例如，环境需要离散动作和连续动作）；
看看如何在策略中使用 RNN（教程）；
将其与决策 Transformers 示例中使用 Transformers 进行比较（请参见 GitHub 上的 example 目录）。

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

估计内存使用量：316 MB

Sphinx-Gallery 生成的图库