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)	接受 `TensorSpec` 作為引數來控制輸出域的 `tensordict.nn.TensorDictModule` 子類。
SafeSequential(*args, **kwargs)	TensorDictModule 的安全序列。
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 還是常規 Actor 取決於實現的演算法。線上演算法通常需要機率 Actor，離線演算法通常有一個確定性 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, ..., **kwargs)	接受 `TensorSpec` 作為引數來控制輸出域的 `tensordict.nn.ProbabilisticTensorDictSequential` 子類。

Q 值 Actor

Q 值 Actor 是一種策略，它根據狀態-動作對的最大值（或“質量”）來選擇動作。這個值可以表示為表格或函式。對於具有連續狀態的離散動作空間，通常使用神經網路等非線性模型來表示這個函式。

QValueActor

QValueActor 類接受一個模組和一個動作規範，並輸出選定的動作及其對應的 Q 值。

>>> 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 學習是 Q 學習的一個變體，它將值函式表示為值可能性的機率分佈，而不是單個標量值。這使得代理能夠了解環境中的不確定性並做出更明智的決策。在 TorchRL 中，分佈 Q 學習使用 `DistributionalQValueActor` 類來實現。該類接受一個模組、一個動作規範和一個支援向量，並輸出選定的動作及其對應的 Q 值分佈。

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

這將輸出一個包含選定動作及其對應 Q 值分佈的張量字典。

QValueActor(*args, **kwargs)	Q 值 Actor 類。
QValueModule(*args, **kwargs)	用於 Q 值策略的 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”），具體取決於輸入。因此，該類同時處理價值網路和 Q 值網路。還提供了三個類來組合策略和價值網路。`ActorCriticOperator` 是一個具有共享引數的聯合 Actor-Critic 網路：它讀取一個觀測，透過一個公共骨幹網路進行傳遞，寫入一個隱藏狀態，將此隱藏狀態饋送給策略，然後採用隱藏狀態和動作，並提供狀態-動作對的 Q 值。`ActorValueOperator` 是一個具有共享引數的聯合 Actor-Value 網路：它讀取一個觀測，透過一個公共骨幹網路進行傳遞，寫入一個隱藏狀態，將此隱藏狀態饋送給策略和價值模組，以輸出動作和狀態值。最後，`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-Value 運算子。
ActorValueOperator(*args, **kwargs)	Actor-Value 運算子。
ValueOperator(*args, **kwargs)	RL 中價值函式的通用類。
DecisionTransformerInferenceWrapper(*args, ..., **kwargs)	Decision Transformer 的推理動作包裝器。

特定領域的 TensorDict 模組

這些模組包括 MBRL 或 RLHF 管道的專用解決方案。

LMHeadActorValueOperator(*args, **kwargs)	從 huggingface 風格的 `LMHeadModel` 構建 Actor-Value 運算子。
WorldModelWrapper(*args, **kwargs)	世界模型包裝器。

Hooks

`QValueActor` 和 `DistributionalQValueActor` 模組使用 Q 值鉤子，通常應優先使用它們，因為它們更容易建立和使用。

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

模型

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

常規模組

BatchRenorm1d(num_features, *, momentum=0.1, **kwargs)	BatchRenorm 模組 (https://arxiv.org/abs/1702.03275)。
ConsistentDropout([p])	實現了具有一致性 dropout 的 `Dropout` 變體。
Conv3dNet(in_features, depth, num_cells, **kwargs)	3D 卷積神經網路。
ConvNet(in_features, depth, num_cells, **kwargs)	卷積神經網路。
MLP(in_features, out_features, depth, **kwargs)	多層感知機。
Squeeze2dLayer()	卷積神經網路的擠壓層。
SqueezeLayer([dims])	擠壓層。

特定於演算法的模組

這些網路實現了對特定演算法（如 DQN、DDPG 或 Dreamer）有用的子網路。

DTActor(state_dim, action_dim[, ...])	Decision 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[, ...])	線上 Decision 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[, ...])	線上 Decision 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, **kwargs)	多智慧體網路的基類。
MultiAgentMLP(n_agent_inputs, ..., **kwargs)	多智慧體 MLP。
MultiAgentConvNet(n_agents, centralized, ..., **kwargs)	多智慧體 CNN。
QMixer(state_shape, mixing_embed_dim, ..., **kwargs)	QMix 混合器。
VDNMixer(n_agents, device)	值分解網路 (VDN) 混合器。

探索

帶噪聲的線性層是探索環境的一種流行方式，它不改變動作，而是將隨機性整合到權重配置中。

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, ...])	實現帶位置縮放的獨立正態分佈。
TanhNormal(loc, scale[, upscale, low, high, ...])	實現帶位置縮放的 Tanh 正態分佈。
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 版本。

Utils

模組實用函式包括用於執行自定義對映的功能以及一個用於從給定模組構建 `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 包裝器。