torchrl.modules 包¶

TensorDict 模块：Actor、探索、价值模型和生成模型¶

TorchRL 提供了一系列模块包装器，旨在简化从头开始构建 RL 模型的过程。这些包装器完全基于 tensordict.nn.TensorDictModule 和 tensordict.nn.TensorDictSequential。它们可以大致分为三类：策略（Actor），包括探索策略、价值模型和模拟模型（在基于模型的上下文中）。

主要特点是

将规范集成到您的模型中，以确保模型输出与您的环境期望的输入相匹配；
概率模块，可以自动从选定的分布中采样和/或返回感兴趣的分布；
用于 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 方式）到期望的域中。

`Actor`(args, *kwargs)	RL 中确定性 Actor 的通用类。
`MultiStepActorWrapper`(args, *kwargs)	多动作 Actor 的包装器。
`SafeModule`(args, *kwargs)	`tensordict.nn.TensorDictModule` 子类，它接受 `TensorSpec` 作为参数来控制输出域。
`SafeSequential`(args, *kwargs)	TensorDictModules 的安全序列。
`TanhModule`(args, *kwargs)	用于具有有界动作空间的确定性策略的 Tanh 模块。

探索包装器和模块¶

为了有效地探索环境，TorchRL 提出了一系列模块，这些模块将用更嘈杂的版本覆盖策略采样的动作。它们的行为由 exploration_type() 控制：如果探索设置为 ExplorationType.RANDOM，则探索处于活动状态。在所有其他情况下，写入 tensordict 的动作只是网络输出。

注意

与其他探索模块不同，ConsistentDropoutModule 使用 train/eval 模式来符合 PyTorch 中的常规 Dropout API。set_exploration_type() 上下文管理器对此模块没有影响。

`AdditiveGaussianModule`(args, *kwargs)	加性高斯 PO 模块。
`AdditiveGaussianWrapper`(args, *kwargs)	加性高斯 PO 包装器。
`ConsistentDropoutModule`(args, *kwargs)	`ConsistentDropout` 的 TensorDictModule 包装器。
`EGreedyModule`(args, *kwargs)	Epsilon-Greedy 探索模块。
`EGreedyWrapper`(args, *kwargs)	[已弃用] Epsilon-Greedy PO 包装器。
`OrnsteinUhlenbeckProcessModule`(args, *kwargs)	Ornstein-Uhlenbeck 探索策略模块。
`OrnsteinUhlenbeckProcessWrapper`(args, *kwargs)	Ornstein-Uhlenbeck 探索策略包装器。

概率 Actor¶

诸如 PPO 之类的某些算法需要实现概率策略。在 TorchRL 中，这些策略采用模型的形式，后跟分布构造器。

注意

概率或常规 Actor 类的选择取决于正在实现的算法。在线策略算法通常需要概率 Actor，离线策略通常具有确定性 Actor 和额外的探索策略。但是，此规则也有许多例外。

该模型读取输入（通常是来自环境的一些观察）并输出分布的参数，而分布构造器读取这些参数并从分布中获取随机样本和/或提供 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" 以及此动作的对数概率。

`ProbabilisticActor`(args, *kwargs)	RL 中概率 Actor 的通用类。
`SafeProbabilisticModule`(args, *kwargs)	`tensordict.nn.ProbabilisticTensorDictModule` 子类，它接受 `TensorSpec` 作为参数来控制输出域。
`SafeProbabilisticTensorDictSequential`(*args, ...)	`tensordict.nn.ProbabilisticTensorDictSequential` 子类，它接受 `TensorSpec` 作为参数来控制输出域。

Q 值 Actor¶

Q 值 Actor 是一种特殊的策略类型，它不直接从观察中预测动作，而是选择最大化 (s,a) -> v 映射的值（或质量）的动作。此映射可以是表或函数。对于具有连续（或接近连续，例如像素）状态的离散动作空间，习惯上对映射使用非线性模型（例如神经网络）。Q 值网络的语义希望非常简单：我们只需要输入一个张量到张量的映射，给定某个状态（输入张量），输出要从中选择的动作值列表。包装器会将生成的动作以及动作值列表写入输入 tensordict 中。

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> 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 = OneHot(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”（箱）。值网络输出状态-动作值属于一个箱或另一个箱的概率。因此，对于维度为 M 的状态空间、维度为 N 的动作空间和箱数 B，值网络编码了 (s,a) -> v 映射。此映射可以是表或函数。对于具有连续（或接近连续，例如像素）状态的离散动作空间，习惯上对映射使用非线性模型（例如神经网络）。Q 值网络的语义希望非常简单：我们只需要输入一个张量到张量的映射，给定某个状态（输入张量），输出要从中选择的动作值列表。包装器会将生成的动作以及动作值列表写入输入 tensordict 中。

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> 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 = OneHot(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”（箱）。值网络输出状态-动作值属于一个箱或另一个箱的概率。因此，对于维度为 M 的状态空间、维度为 N 的动作空间和箱数 B，值网络编码了 (s,a) -> v 映射。此映射可以是表或函数。对于具有连续（或接近连续，例如像素）状态的离散动作空间，习惯上对映射使用非线性模型（例如神经网络）。Q 值网络的语义希望非常简单：我们只需要输入一个张量到张量的映射，给定某个状态（输入张量），输出要从中选择的动作值列表。包装器会将生成的动作以及动作值列表写入输入 tensordict 中。

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> 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 = OneHot(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”（箱）。值网络输出状态-动作值属于一个箱或另一个箱的概率。因此，对于维度为 M 的状态空间、维度为 N 的动作空间和箱数 B，值网络编码了 \(\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 OneHot
>>> 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 = OneHot(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)

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> 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 = OneHot(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)

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> 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 = OneHot(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)

`QValueActor`(args, *kwargs)	Q 值 Actor 类。
`QValueModule`(args, *kwargs)	用于 Q 值策略的 Q 值 TensorDictModule。
`DistributionalQValueActor`(args, *kwargs)	分布 DQN Actor 类。
`DistributionalQValueModule`(args, *kwargs)	用于 Q 值策略的分布 Q 值钩子。

价值运算符和联合模型¶

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

`ActorCriticOperator`(args, *kwargs)	Actor-critic 运算符。
`ActorCriticWrapper`(args, *kwargs)	没有公共模块的 Actor-价值运算符。
`ActorValueOperator`(args, *kwargs)	Actor-价值运算符。
`ValueOperator`(args, *kwargs)	RL 中价值函数的通用类。
`DecisionTransformerInferenceWrapper`(*args, ...)	决策 Transformer 的推理动作包装器。

特定领域的 TensorDict 模块¶

这些模块包括用于 MBRL 或 RLHF 管道的专用解决方案。

`LMHeadActorValueOperator`(args, *kwargs)	从类似 huggingface 的 *LMHeadModel 构建 Actor-Value 运算符。
`WorldModelWrapper`(args, *kwargs)	世界模型包装器。

钩子¶

Q 值钩子由 QValueActor 和 DistributionalQValueActor 模块使用，通常应优先使用这些模块，因为它们更易于创建和使用。

`QValueHook`(action_space[, var_nums, ...])	用于 Q 值策略的 Q 值钩子。
`DistributionalQValueHook`(action_space, support)	用于 Q 值策略的分布 Q 值钩子。

模型¶

TorchRL 提供了一系列有用的“常规”（即非 tensordict）nn.Module 类，用于 RL 用途。

常规模块¶

`BatchRenorm1d`(num_features, *[, momentum, ...])	BatchRenorm 模块 (https://arxiv.org/abs/1702.03275)。
`Conv3dNet`(in_features, depth, num_cells, ...)	一个 3D 卷积神经网络。
`ConvNet`(in_features, depth, num_cells, ...)	一个卷积神经网络。
`MLP`(in_features, out_features, ...)	一个多层感知器。
`Squeeze2dLayer`()	用于卷积神经网络的挤压层。
`SqueezeLayer`([dims])	挤压层。

特定于算法的模块¶

这些网络实现了已被证明对特定算法（例如 DQN、DDPG 或 Dreamer）有用的子网络。

`DTActor`(state_dim, action_dim[, ...])	决策 Transformer Actor 类。
`DdpgCnnActor`(action_dim[, conv_net_kwargs, ...])	DDPG 卷积 Actor 类。
`DdpgCnnQNet`([conv_net_kwargs, ...])	DDPG 卷积 Q 值类。
`DdpgMlpActor`(action_dim[, mlp_net_kwargs, ...])	DDPG Actor 类。
`DdpgMlpQNet`([mlp_net_kwargs_net1, ...])	DDPG Q 值 MLP 类。
`DecisionTransformer`(state_dim, action_dim[, ...])	在线决策 Transformer。
`DistributionalDQNnet`(args, *kwargs)	分布深度 Q 网络 softmax 层。
`DreamerActor`(out_features[, depth, ...])	Dreamer Actor 网络。
`DuelingCnnDQNet`(out_features[, ...])	Dueling CNN Q 网络。
`GRUCell`(input_size, hidden_size[, bias, ...])	门控循环单元 (GRU) 单元，其执行与 nn.LSTMCell 相同的操作，但完全以 Python 编码。
`GRU`(input_size, hidden_size[, num_layers, ...])	用于执行多层 GRU 的多个步骤的 PyTorch 模块。
`GRUModule`(args, *kwargs)	GRU 模块的嵌入器。
`LSTMCell`(input_size, hidden_size[, bias, ...])	长短期记忆 (LSTM) 单元，其执行与 nn.LSTMCell 相同的操作，但完全以 Python 编码。
`LSTM`(input_size, hidden_size[, num_layers, ...])	用于执行多层 LSTM 的多个步骤的 PyTorch 模块。
`LSTMModule`(args, *kwargs)	LSTM 模块的嵌入器。
`ObsDecoder`([channels, num_layers, ...])	观察解码器网络。
`ObsEncoder`([channels, num_layers, depth])	观察编码器网络。
`OnlineDTActor`(state_dim, action_dim[, ...])	在线决策 Transformer Actor 类。
`RSSMPosterior`([hidden_dim, state_dim, scale_lb])	RSSM 的后验网络。
`RSSMPrior`(action_spec[, hidden_dim, ...])	RSSM 的先验网络。

多智能体特定模块¶

这些网络实现了可用于多智能体上下文的模型。它们使用 vmap() 一次在网络输入上执行多个网络。由于参数是批量处理的，因此初始化可能与通常使用其他 PyTorch 模块的方式不同，请参阅 get_stateful_net() 以获取更多信息。

`MultiAgentNetBase`(*, n_agents[, ...])	多智能体网络的基础类。
`MultiAgentMLP`(n_agent_inputs, ...)	多智能体 MLP。
`MultiAgentConvNet`(n_agents, centralized, ...)	多智能体 CNN。
`QMixer`(state_shape, mixing_embed_dim, ...)	QMix 混合器。
`VDNMixer`(n_agents, device)	价值分解网络混合器。

探索¶

噪声线性层是一种流行的探索环境的方式，它不改变动作，而是将随机性集成到权重配置中。

`NoisyLinear`(in_features, out_features[, ...])	噪声线性层。
`NoisyLazyLinear`(out_features[, bias, ...])	噪声惰性线性层。
`reset_noise`(layer)	重置噪声层的噪声。

规划器¶

`CEMPlanner`(args, *kwargs)	CEMPlanner 模块。
`MPCPlannerBase`(args, *kwargs)	MPCPlannerBase 抽象模块。
`MPPIPlanner`(args, *kwargs)	MPPI Planner 模块。

分布¶

一些分布通常在 RL 脚本中使用。

`Delta`(param[, atol, rtol, batch_shape, ...])	Delta 分布。
`IndependentNormal`(loc, scale[, upscale, ...])	实现具有位置缩放的 Normal 分布。
`NormalParamWrapper`(operator[, ...])	用于 normal 分布参数的包装器。
`TanhNormal`(loc, scale[, upscale, low, high, ...])	实现具有位置缩放的 TanhNormal 分布。
`TruncatedNormal`(loc, scale[, upscale, low, ...])	实现具有位置缩放的截断 Normal 分布。
`TanhDelta`(param[, low, high, event_dims, ...])	实现 Tanh 变换的 Delta 分布。
`OneHotCategorical`([logits, probs, grad_method])	One-hot 分类分布。
`MaskedCategorical`([logits, probs, mask, ...])	MaskedCategorical 分布。
`MaskedOneHotCategorical`([logits, probs, ...])	MaskedCategorical 分布。

实用工具¶

`mappings`(key)	给定一个输入字符串，返回一个满射函数 f(x): R -> R^+。
`inv_softplus`(bias)	反向 softplus 函数。
`biased_softplus`(bias[, min_val])	一个有偏置的 softplus 模块。
`get_primers_from_module`(module)	从模块的所有子模块中获取所有 tensordict primers。

VmapModule(*args, **kwargs)

一个 TensorDictModule 包装器，用于在输入上进行 vmap 操作。