torchrl.modules 包¶

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

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

主要特性是：

将 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 中确定性 Actor 的通用类。
`MultiStepActorWrapper`(args, *kwargs)	多动作 Actor 的包装器。
`SafeModule`(args, *kwargs)	接受 `TensorSpec` 作为参数以控制输出域的 `tensordict.nn.TensorDictModule` 子类。
`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 探索策略模块。

概率 Actor¶

某些算法（如 PPO）要求实现概率策略。在 TorchRL 中，这些策略的形式是模型，后面跟着一个分布构造器。

注意

选择概率或常规 Actor 类取决于正在实现的算法。On-policy 算法通常需要概率 Actor，off-policy 通常具有具有额外探索策略的确定性 Actor。但是，这个规则有很多例外。

模型读取输入（通常是环境的某个观察值）并输出分布的参数，而分布构造器读取这些参数并获取分布的随机样本和/或提供一个 `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 中概率 Actor 的通用类。
`SafeProbabilisticModule`(args, *kwargs)	接受 `TensorSpec` 作为参数以控制输出域的 `tensordict.nn.ProbabilisticTensorDictModule` 子类。
`SafeProbabilisticTensorDictSequential`(*args, ...)	接受 `TensorSpec` 作为参数以控制输出域的 `tensordict.nn.ProbabilisticTensorDictSequential` 子类。

Q 值 Actor¶

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

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)

这将输出一个包含选定动作及其对应值的张量字典。

分布 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)

这将输出一个包含选定动作及其对应值分布的张量字典。

`QValueActor`(args, *kwargs)	Q 值 Actor 类。
`QValueModule`(args, *kwargs)	Q 值 Actor 策略的 Q 值 TensorDictModule。
`DistributionalQValueActor`(args, *kwargs)	分布 DQN Actor 类。
`DistributionalQValueModule`(args, *kwargs)	Q 值 Actor 策略的分布 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

`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)	世界模型包装器。

Hooks¶

`QValueActor` 和 `DistributionalQValueActor` 模块使用 Q 值 Hook，并且通常应该优先使用它们，因为它们更易于创建和使用。

`QValueHook`(action_space[, var_nums, ...])	Q 值 Actor 策略的 Q 值 Hook。
`DistributionalQValueHook`(action_space, support)	Q 值 Actor 策略的分布 Q 值 Hook。

模型¶

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

常规模块¶

`BatchRenorm1d`(num_features, *[, momentum, ...])	BatchRenorm 模块（https://arxiv.org/abs/1702.03275）。
`ConsistentDropout`([p])	实现了一个具有一致 Dropout 的 Dropout 变体（`Dropout`）。
`Conv3dNet`(in_features, depth, num_cells, ...)	3D 卷积神经网络。
`ConvNet`(in_features, depth, num_cells, ...)	卷积神经网络。
`MLP`(in_features, out_features, depth, ...)	多层感知机。
`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, ...])	执行与 `nn.LSTMCell` 相同操作的门控循环单元（GRU）单元，但完全用 Python 编写。
`GRU`(input_size, hidden_size[, num_layers, ...])	用于执行多层 GRU 多步的 PyTorch 模块。
`GRUModule`(args, *kwargs)	GRU 模块的嵌入器。
`LSTMCell`(input_size, hidden_size[, bias, ...])	执行与 `nn.LSTMCell` 相同操作的长短期记忆（LSTM）单元，但完全用 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 的先验网络。
`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 规划器模块。

分布¶

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

`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）分类分布。
`LLMMaskedCategorical`(logits, mask[, ...])	为 LLM 优化的掩码分类分布。
`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)

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