torchrl.modules 包¶

TensorDict 模块: 智能体、探索、值模型和生成模型¶

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

主要功能包括

将 spec 集成到模型中，确保模型输出与环境期望的输入匹配；
概率模块，可以自动从选定的分布中采样和/或返回感兴趣的分布；
用于 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 中确定性智能体的通用类。
`MultiStepActorWrapper`(args, *kwargs)	多动作智能体的包装器。
`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 模块。
`ConsistentDropoutModule`(args, *kwargs)	`ConsistentDropout` 的 TensorDictModule 包装器。
`EGreedyModule`(args, *kwargs)	Epsilon-Greedy 探索模块。
`OrnsteinUhlenbeckProcessModule`(args, *kwargs)	Ornstein-Uhlenbeck 探索策略模块。

概率智能体¶

PPO 等一些算法需要实现概率策略。在 TorchRL 中，这些策略的形式是模型，后跟一个分布构造器。

注意

选择概率或常规智能体类取决于正在实现的算法。在策略算法通常需要概率智能体，离策略算法通常具有带额外探索策略的确定性智能体。然而，这条规则也有许多例外。

模型读取输入（通常是来自环境的一些观察），并输出分布的参数，而分布构造器读取这些参数并从分布中获取随机样本和/或提供一个 torch.distributions.Distribution 对象。

>>> from tensordict.nn import NormalParamExtractor, TensorDictSequential, TensorDictModule
>>> from torchrl.modules import SafeProbabilisticModule
>>> from torchrl.envs import GymEnv
>>> from torch.distributions import Normal
>>> from torch import nn
>>>
>>> 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

>>> from torchrl.modules import 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 中概率智能体的通用类。
`SafeProbabilisticModule`(args, *kwargs)	`tensordict.nn.ProbabilisticTensorDictModule` 的子类，接受 `TensorSpec` 作为参数来控制输出域。
`SafeProbabilisticTensorDictSequential`(*args, ...)	`tensordict.nn.ProbabilisticTensorDictSequential` 的子类，接受 `TensorSpec` 作为参数来控制输出域。

Q 值智能体¶

Q 值智能体是一种策略类型，它根据状态-动作对的最大值（或“质量”）来选择动作。这个值可以表示为一个表格或一个函数。对于具有连续状态的离散动作空间，通常使用神经网络等非线性模型来表示此函数。

QValueActor¶

QValueActor 类接受一个模块和一个动作规范，并输出选定的动作及其对应的值。

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> from torchrl.modules.tensordict_module.actors import QValueActor
>>> # Create a tensor dict with an observation
>>> td = TensorDict({'observation': torch.randn(5, 3)}, [5])
>>> # Define the action space
>>> action_spec = OneHot(4)
>>> # Create a linear module to output action values
>>> module = nn.Linear(3, 4)
>>> # Create a QValueActor instance
>>> qvalue_actor = QValueActor(module=module, spec=action_spec)
>>> # Run the actor on the tensor dict
>>> 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)

这将输出一个 tensor dict，其中包含选定的动作及其对应的值。

分布式 Q 学习¶

分布式 Q 学习是 Q 学习的一种变体，它将值函数表示为可能值的概率分布，而不是单个标量值。这使得智能体能够了解环境中的不确定性，并做出更明智的决策。在 TorchRL 中，分布式 Q 学习使用 DistributionalQValueActor 类实现。此类接受一个模块、一个动作规范和一个支持向量，并输出选定的动作及其对应的值分布。

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from torchrl.data import OneHot
>>> from torchrl.modules import DistributionalQValueActor, MLP
>>> # Create a tensor dict with an observation
>>> td = TensorDict({'observation': torch.randn(5, 4)}, [5])
>>> # Define the action space
>>> action_spec = OneHot(4)
>>> # Define the number of bins for the value distribution
>>> nbins = 3
>>> # Create an MLP module to output logits for the value distribution
>>> module = MLP(out_features=(nbins, 4), depth=2)
>>> # Create a DistributionalQValueActor instance
>>> qvalue_actor = DistributionalQValueActor(module=module, spec=action_spec, support=torch.arange(nbins))
>>> # Run the actor on the tensor dict
>>> 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)

这将输出一个 tensor dict，其中包含选定的动作及其对应的值分布。

`QValueActor`(args, *kwargs)	一个 Q 值智能体类。
`QValueModule`(args, *kwargs)	用于 Q 值策略的 Q 值 TensorDictModule。
`DistributionalQValueActor`(args, *kwargs)	一个分布式 DQN 智能体类。
`DistributionalQValueModule`(args, *kwargs)	用于 Q 值策略的分布式 Q 值 hook。

值算子和联合模型¶

TorchRL 提供了一系列值算子，用于包装值网络以简化与库其余部分的接口。基本构建块是 torchrl.modules.tensordict_module.ValueOperator：给定输入状态（可能还有动作），它将根据输入自动在 tensordict 中写入 "state_value"（或 "state_action_value"）。因此，此类同时处理值网络和质量网络。还提出了三个类来组合策略和值网络。ActorCriticOperator 是一个具有共享参数的联合智能体-质量网络：它读取观察，通过一个共同骨干网络，写入隐藏状态，将此隐藏状态馈送给策略，然后结合隐藏状态和动作，提供状态-动作对的质量。ActorValueOperator 是一个具有共享参数的联合智能体-值网络：它读取观察，通过一个共同骨干网络，写入隐藏状态，将此隐藏状态馈送给策略和值模块，以输出动作和状态值。最后，ActorCriticWrapper 是一个没有共享参数的联合智能体和值网络。它主要用作 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)	智能体-评论家算子。
`ActorCriticWrapper`(args, *kwargs)	没有共同模块的智能体-值算子。
`ActorValueOperator`(args, *kwargs)	智能体-值算子。
`ValueOperator`(args, *kwargs)	RL 中值函数的通用类。
`DecisionTransformerInferenceWrapper`(*args, ...)	用于 Decision Transformer 的推理动作包装器。

领域特定 TensorDict 模块¶

这些模块包括用于 MBRL 或 RLHF 流水线的专用解决方案。

`LMHeadActorValueOperator`(args, *kwargs)	从类似于 huggingface 的 *LMHeadModel 构建智能体-值算子。
`WorldModelWrapper`(args, *kwargs)	世界模型包装器。

Hooks¶

Q 值 hook 被 QValueActor 和 DistributionalQValueActor 模块使用，通常应优先选择这些模块，因为它们更易于创建和使用。

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

模型¶

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

常规模块¶

`BatchRenorm1d`(num_features, *[, momentum, ...])	BatchRenorm 模块 (https://arxiv.org/abs/1702.03275)。
`ConsistentDropout`([p])	实现了具有一致 dropout 的 `Dropout` 变体。
`Conv3dNet`(in_features, depth, num_cells, ...)	一个 3D 卷积神经网络。
`ConvNet`(in_features, depth, num_cells, ...)	一个卷积神经网络。
`MLP`(in_features, out_features, depth, ...)	一个多层感知机。
`Squeeze2dLayer`()	用于卷积神经网络的 Squeezing 层。
`SqueezeLayer`([dims])	Squeezing 层。

算法特定模块¶

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

`DTActor`(state_dim, action_dim[, ...])	Decision Transformer 智能体类。
`DdpgCnnActor`(action_dim[, conv_net_kwargs, ...])	DDPG 卷积智能体类。
`DdpgCnnQNet`([conv_net_kwargs, ...])	DDPG 卷积 Q 值类。
`DdpgMlpActor`(action_dim[, mlp_net_kwargs, ...])	DDPG 智能体类。
`DdpgMlpQNet`([mlp_net_kwargs_net1, ...])	DDPG Q 值 MLP 类。
`DecisionTransformer`(state_dim, action_dim[, ...])	在线 Decision Transformer。
`DistributionalDQNnet`(args, *kwargs)	分布式深度 Q 网络 softmax 层。
`DreamerActor`(out_features[, depth, ...])	Dreamer 智能体网络。
`DuelingCnnDQNet`(out_features[, ...])	Dueling CNN Q 网络。
`GRUCell`(input_size, hidden_size[, bias, ...])	一个门控循环单元 (GRU) cell，执行与 nn.LSTMCell 相同的操作，但完全用 Python 代码实现。
`GRU`(input_size, hidden_size[, num_layers, ...])	一个用于执行多层 GRU 多步操作的 PyTorch 模块。
`GRUModule`(args, *kwargs)	GRU 模块的嵌入器。
`LSTMCell`(input_size, hidden_size[, bias, ...])	一个长短期记忆 (LSTM) cell，执行与 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[, ...])	在线 Decision Transformer 智能体类。
`RSSMPosterior`([hidden_dim, state_dim, scale_lb])	RSSM 的后验网络。
`RSSMPrior`(action_spec[, hidden_dim, ...])	RSSM 的先验网络。
`set_recurrent_mode`([mode])	用于设置 RNN 循环模式的上下文管理器。
`recurrent_mode`()	返回当前的采样类型。

多智能体专用模块¶

这些网络实现了可用于多智能体环境的模型。它们使用 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 规划器模块。

分布¶

某些分布通常用于强化学习脚本。

`Delta`(param[, atol, rtol, batch_shape, ...])	Delta 分布。
`IndependentNormal`(loc, scale[, upscale, ...])	实现具有位置缩放的正态分布。
`TanhNormal`(loc, scale[, upscale, low, high, ...])	实现具有位置缩放的 TanhNormal 分布。
`TruncatedNormal`(loc, scale[, upscale, low, ...])	实现具有位置缩放的截断正态分布。
`TanhDelta`(param[, low, high, event_dims, ...])	实现 Tanh 转换的 Delta 分布。
`OneHotCategorical`([logits, probs, grad_method])	One-hot 分类分布。
`MaskedCategorical`([logits, probs, mask, ...])	MaskedCategorical 分布。
`MaskedOneHotCategorical`([logits, probs, ...])	MaskedCategorical 分布。
`Ordinal`(scores)	一种用于学习从有限有序集合采样的离散分布。
`OneHotOrdinal`(scores)	`Ordinal` 分布的 one-hot 版本。

工具函数¶

模块工具函数包括用于进行自定义映射的功能，以及一个从给定模块构建 TensorDictPrimer 实例的工具。

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

VmapModule(*args, **kwargs)

一个用于对输入进行 vmap 操作的 TensorDictModule 包装器。