torchrl.modules 包¶
TensorDict 模块：演员、探索、价值模型和生成模型¶
TorchRL 提供了一系列模块包装器，旨在使从头开始构建 RL 模型变得容易。这些包装器完全基于 tensordict.nn.TensorDictModule 和 tensordict.nn.TensorDictSequential。它们可以大致分为三类：策略（演员），包括探索策略、价值模型和模拟模型（在基于模型的上下文中）。
主要功能是
将规范集成到您的模型中以确保模型输出与您的环境期望的输入相匹配；
概率模块，可以从选定的分布中自动采样和/或返回感兴趣的分布；
用于 Q 值学习、基于模型的代理和其他方面的自定义容器。
TensorDictModules 和 SafeModules¶
TorchRL SafeModule 允许您检查您的模型输出是否与环境期望的输出相匹配。当您的模型要跨多个环境循环使用时，以及当您希望确保输出（例如，动作）始终满足环境施加的界限时，应使用此功能。以下是如何使用该功能与 Actor 类一起使用的示例
>>> env = GymEnv("Pendulum-v1")
>>> action_spec = env.action_spec
>>> model = nn.LazyLinear(action_spec.shape[-1])
>>> policy = Actor(model, in_keys=["observation"], spec=action_spec, safe=True)
safe 标志确保输出始终在 action_spec 域的界限内：如果网络输出违反这些界限，它将被投影（以 L1 方式）到所需的域。
|
RL 中确定性演员的通用类。 |
|
多动作演员的包装器。 |
|
|
|
TensorDictModules 的安全序列。 |
|
用于具有有界动作空间的确定性策略的 Tanh 模块。 |
探索包装器¶
为了有效地探索环境，TorchRL 提出了 一系列包装器，这些包装器将通过更嘈杂的版本覆盖策略采样的动作。它们的 行为由 exploration_mode() 控制：如果探索设置为 "random"，则探索处于活动状态。在所有其他情况下，tensordict 中写入的动作只是网络输出。
|
加性高斯 PO 模块。 |
|
加性高斯 PO 包装器。 |
|
Epsilon-Greedy 探索模块。 |
|
[已弃用] Epsilon-Greedy PO 包装器。 |
|
Ornstein-Uhlenbeck 探索策略模块。 |
|
Ornstein-Uhlenbeck 探索策略包装器。 |
概率性演员¶
某些算法（例如 PPO）需要实现概率性策略。在 TorchRL 中，这些策略采用模型形式，然后是分布构造函数。
注意
概率性演员类或常规演员类取决于正在实现的算法。在线策略算法通常需要概率性演员，离线策略通常具有确定性演员，并具有额外的探索策略。但是，此规则也有许多例外情况。
模型读取输入（通常是来自环境的一些观察结果）并输出分布的参数，而分布构造函数读取这些参数并从分布中获取随机样本和/或提供 torch.distributions.Distribution 对象。
>>> from tensordict.nn import NormalParamExtractor, TensorDictSequential
>>> from torch.distributions import Normal
>>> env = GymEnv("Pendulum-v1")
>>> action_spec = env.action_spec
>>> model = nn.Sequential(nn.LazyLinear(action_spec.shape[-1] * 2), NormalParamExtractor())
>>> # build the first module, which maps the observation on the mean and sd of the normal distribution
>>> model = TensorDictModule(model, in_keys=["observation"], out_keys=["loc", "scale"])
>>> # build the distribution constructor
>>> prob_module = SafeProbabilisticModule(
... in_keys=["loc", "scale"],
... out_keys=["action"],
... distribution_class=Normal,
... return_log_prob=True,
... spec=action_spec,
... )
>>> policy = TensorDictSequential(model, prob_module)
>>> # execute a rollout
>>> env.rollout(3, policy)
为了便于构建概率性策略，我们提供了一个专用的 ProbabilisticActor
>>> policy = ProbabilisticActor(
... model,
... in_keys=["loc", "scale"],
... out_keys=["action"],
... distribution_class=Normal,
... return_log_prob=True,
... spec=action_spec,
... )
这消除了指定构造函数并将构造函数与模块按顺序放置在一起的需要。
此策略的输出将包含 "loc" 和 "scale" 条目，一个根据正态分布采样的 "action" 以及此动作的对数概率。
|
RL 中概率 actor 的通用类。 |
|
|
|
|
Q 值 actor¶
Q 值 actor 是一种特殊的策略类型，它不直接根据观察预测动作，而是选择使 (s,a) -> v 映射的值（或质量）最大化的动作。此映射可以是表格或函数。对于具有连续（或接近连续，例如像素）状态的离散动作空间，习惯上使用诸如神经网络之类的非线性模型来进行映射。Q 值网络的语义希望非常简单：我们只需要提供一个张量到张量映射，该映射在给定特定状态（输入张量）的情况下，输出可以选择的动作值的列表。包装器将根据输入 tensordict 中的动作值列表写入生成的 action。
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch import nn
>>> from torchrl.data import OneHotDiscreteTensorSpec
>>> from torchrl.modules.tensordict_module.actors import QValueActor
>>> td = TensorDict({'observation': torch.randn(5, 3)}, [5])
>>> # we have 4 actions to choose from
>>> action_spec = OneHotDiscreteTensorSpec(4)
>>> # the model reads a state of dimension 3 and outputs 4 values, one for each action available
>>> module = nn.Linear(3, 4)
>>> qvalue_actor = QValueActor(module=module, spec=action_spec)
>>> qvalue_actor(td)
>>> print(td)
TensorDict(
fields={
action: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.int64, is_shared=False),
action_value: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
chosen_action_value: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
observation: Tensor(shape=torch.Size([5, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([5]),
device=None,
is_shared=False)
分布式 Q 学习略有不同：在这种情况下，值网络不会为每个状态动作值输出标量值。相反，值空间被划分为任意数量的“bins”。值网络输出状态动作值属于一个或另一个 bin 的概率。因此，对于维数为 M 的状态空间、维数为 N 的动作空间以及数量为 B 的 bins，值网络编码一个 \(\mathbb{R}^{M} \rightarrow \mathbb{R}^{N \times B}\) 映射。以下示例演示了如何在 TorchRL 中使用 DistributionalQValueActor 类来实现这一点。
>>> import torch
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from torchrl.data import OneHotDiscreteTensorSpec
>>> from torchrl.modules import DistributionalQValueActor, MLP
>>> td = TensorDict({'observation': torch.randn(5, 4)}, [5])
>>> nbins = 3
>>> # our model reads the observation and outputs a stack of 4 logits (one for each action) of size nbins=3
>>> module = MLP(out_features=(nbins, 4), depth=2)
>>> action_spec = OneHotDiscreteTensorSpec(4)
>>> qvalue_actor = DistributionalQValueActor(module=module, spec=action_spec, support=torch.arange(nbins))
>>> td = qvalue_actor(td)
>>> print(td)
TensorDict(
fields={
action: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.int64, is_shared=False),
action_value: Tensor(shape=torch.Size([5, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
observation: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
batch_size=torch.Size([5]),
device=None,
is_shared=False)
|
Q 值 actor 类。 |
|
用于 Q 值策略的 Q 值 TensorDictModule。 |
|
分布式 DQN actor 类。 |
|
用于 Q 值策略的分布式 Q 值 hook。 |
值运算符和组合模型¶
TorchRL 提供了一系列值运算符，这些运算符包装值网络以简化与库其余部分的接口。基本构建块是 torchrl.modules.tensordict_module.ValueOperator：给定输入状态（可能还有动作），它将自动在 tensordict 中写入 "state_value"（或 "state_action_value"），具体取决于输入内容。因此，此类同时考虑了值网络和质量网络。还提出了三个类来将策略和值网络分组在一起。 ActorCriticOperator 是一个具有共享参数的组合 actor-质量网络：它读取观察结果，将其传递到公共骨干，写入隐藏状态，将此隐藏状态馈送到策略，然后获取隐藏状态和动作并提供状态-动作对的质量。 ActorValueOperator 是一个具有共享参数的组合 actor-值网络：它读取观察结果，将其传递到公共骨干，写入隐藏状态，将此隐藏状态馈送到策略和值模块以输出动作和状态值。最后， ActorCriticWrapper 是一个不共享参数的组合 actor 和值网络。它主要用作 ActorValueOperator 的替代品，当脚本需要考虑这两种选择时。
>>> actor = make_actor()
>>> value = make_value()
>>> if shared_params:
... common = make_common()
... model = ActorValueOperator(common, actor, value)
... else:
... model = ActorValueOperator(actor, value)
>>> policy = model.get_policy_operator() # will work in both cases
|
Actor-critic 运算符。 |
|
没有公共模块的 Actor-value 运算符。 |
|
Actor-value 运算符。 |
|
RL 中值函数的通用类。 |
|
Decision Transformer 的推断动作包装器。 |
特定于领域的 TensorDict 模块¶
这些模块包括针对 MBRL 或 RLHF 管道提供的专用解决方案。
|
从类似 huggingface 的 *LMHeadModel 构建 Actor-Value 运算符。 |
|
世界模型包装器。 |
hook¶
Q 值 hook 由 QValueActor 和 DistributionalQValueActor 模块使用，通常应该优先使用它们，因为它们更容易创建和使用。
|
用于 Q 值策略的 Q 值 hook。 |
|
用于 Q 值策略的分布式 Q 值 hook。 |
模型¶
TorchRL 提供了一系列用于 RL 用途的有用“常规”（即非 tensordict）nn.Module 类。
常规模块¶
|
多层感知器。 |
|
卷积神经网络。 |
|
三维卷积神经网络。 |
|
压缩层。 |
用于卷积神经网络的压缩层。 |
特定于算法的模块¶
这些网络实现了已被证明对特定算法（例如 DQN、DDPG 或 Dreamer）有用的子网络。
|
Decision Transformer Actor 类。 |
|
DDPG 卷积 Actor 类。 |
|
DDPG 卷积 Q 值类。 |
|
DDPG Actor 类。 |
|
DDPG Q 值 MLP 类。 |
|
在线决策 Transformer。 |
|
分布式 Deep Q-Network softmax 层。 |
|
Dreamer 演员网络。 |
|
Dueling CNN Q 网络。 |
|
一个门控循环单元 (GRU) 单元，它执行与 nn.LSTMCell 相同的操作，但完全用 Python 编码。 |
|
一个用于执行多层 GRU 多个步骤的 PyTorch 模块。 |
|
一个 GRU 模块的嵌入器。 |
|
一个长短期记忆 (LSTM) 单元，它执行与 nn.LSTMCell 相同的操作，但完全用 Python 编码。 |
|
一个用于执行多层 LSTM 多个步骤的 PyTorch 模块。 |
|
一个 LSTM 模块的嵌入器。 |
|
观察解码器网络。 |
|
观察编码器网络。 |
|
在线决策 Transformer Actor 类。 |
|
RSSM 的后验网络。 |
|
RSSM 的先验网络。 |
多智能体特定模块¶
这些网络实现可在多智能体环境中使用的模型。它们使用 vmap() 在网络输入上一次执行多个网络。由于参数是批处理的，初始化可能与通常使用其他 PyTorch 模块时不同，有关更多信息，请参见 get_stateful_net()。
|
多智能体网络的基类。 |
|
多智能体 MLP。 |
|
多智能体 CNN。 |
|
QMix 混合器。 |
|
值分解网络混合器。 |
探索¶
噪声线性层是一种流行的探索环境的方法，无需改变动作，而是通过在权重配置中集成随机性。
|
噪声线性层。 |
|
噪声延迟线性层。 |
|
重置噪声层的噪声。 |
规划器¶
|
CEMPlanner 模块。 |
|
MPCPlannerBase 抽象模块。 |
|
MPPI 规划器模块。 |
分布¶
一些分布通常用于 RL 脚本。
|
Delta 分布。 |
|
实现具有位置缩放的正态分布。 |
|
正态分布参数的包装器。 |
|
实现具有位置缩放的 TanhNormal 分布。 |
|
实现具有位置缩放的截断正态分布。 |
|
实现 Tanh 转换的 Delta 分布。 |
|
独热分类分布。 |
|
MaskedCategorical 分布。 |
|
MaskedCategorical 分布。 |
工具¶
|
给定一个输入字符串，返回一个满射函数 f(x): R -> R^+。 |
|
逆 softplus 函数。 |
|
一个有偏 softplus 模块。 |
|
从模块的所有子模块中获取所有 tensordict 引导程序。 |
|
一个 TensorDictModule 包装器，用于在输入上进行 vmap。 |
