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 包裝器。