目录
目录

注意

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

开始使用TorchRL的模块

作者: Vincent Moens

注意

如需在笔记本中运行本教程，请在开头添加一个安装单元，内容为：

!pip install tensordict
!pip install torchrl

强化学习旨在构建能够有效应对特定任务的策略。策略可以呈现多种形式，既可以是将观测空间映射到动作空间的可微函数，也可以是更为临时性的方法（例如，对每个可能动作所计算出的一组值执行 argmax 操作）。策略可以是确定性的，也可以是随机性的，并且可能包含复杂组件，例如循环神经网络（RNN）或 Transformer。

兼顾所有这些场景可能相当复杂。在本简明教程中，我们将深入探讨 TorchRL 在策略构建方面的核心功能。我们主要聚焦于两种常见场景下的随机策略和 Q 值策略：分别以多层感知机（MLP）和卷积神经网络（CNN）作为主干网络。

TensorDictModules

与环境和 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"],
)

至此，我们的策略即可执行！使用延迟模块（lazy module） 使我们无需手动获取观测空间（observation space）的形状，因为该模块会自动确定其形状。 现在，该策略已可直接在环境中运行：

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)

可用的专用 TensorDictModule 列表请参见 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通过将各种操作（如从参数构建分布、从该分布采样以及检索对数概率）组合在一个类下来简化此类模块的设计。在这里，我们将使用三个组件构建一个依赖于常规正态分布的演员：

  • 一个 MLP 主干网络读取大小为 [3] 的观测值并输出一个大小为 [2] 的单个张量；

  • 一个NormalParamExtractor模块，它 将把这个输出分成两部分，一个是均值，另一个是标准差，大小为[1];

  • 一个 ProbabilisticActor 将会 读取这些参数作为 in_keys，创建一个分布并将它们填充到我们的张量字典中，包括样本和对数概率。

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)

关于这次发布，有几点需要注意：

  • 由于我们在构建执行器（actor）时提出了该请求，因此还会记录在当时分布下所采取动作的对数概率。这对于 PPO 等算法是必需的。

  • 分布的参数也将在输出张量字典中返回，位于"loc"和"scale"条目下。

您可以控制动作的采样，以使用分布的期望值或其他属性而不是随机样本，如果您的应用程序需要的话。这可以通过 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]))

我们的 \(\epsilon\)-贪婪探索模块通常会根据一定数量的退火帧和初始值进行自定义 \(\epsilon\) 参数。 \(\epsilon = 1\) 的值意味着每个动作都是随机的，而 \(\epsilon=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-Value 演员

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

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

OneHot(
    shape=torch.Size([2]),
    space=CategoricalBox(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的Decision Transformers示例进行比较（请参阅GitHub上的example目录）。

脚本总运行时间： (0 分钟 48.412 秒)

估计内存使用量： 315 MB

通过 Sphinx-Gallery 生成的画廊

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源