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)	一個安全的 TensorDictModule 序列。
`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)

Distributional Q-learning 略有不同：在這種情況下，價值網路 (value network) 不會為每個狀態-動作價值輸出一個純量值。相反地，價值空間被劃分為任意數量的「區間 (bins)」。價值網路輸出狀態-動作價值屬於某個區間的機率。因此，對於維度為 M 的狀態空間、維度為 N 的動作空間以及區間數 B，價值網路會編碼一個 (s,a) -> v 的映射。這個映射可以是一個表格或一個函數。對於具有連續（或接近連續，例如像素）狀態的離散動作空間，習慣上使用非線性模型（例如神經網路）進行映射。Q-Value 網路的語義 (semantic) 希望非常簡單：我們只需要提供一個 tensor-to-tensor 映射，給定一個特定的狀態（輸入 tensor），輸出一個動作價值列表以供選擇。封裝器 (wrapper) 會將產生的動作連同動作價值列表寫入輸入的 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)

Distributional Q-learning 略有不同：在這種情況下，價值網路 (value network) 不會為每個狀態-動作價值輸出一個純量值。相反地，價值空間被劃分為任意數量的「區間 (bins)」。價值網路輸出狀態-動作價值屬於某個區間的機率。因此，對於維度為 M 的狀態空間、維度為 N 的動作空間以及區間數 B，價值網路會編碼一個 (s,a) -> v 的映射。這個映射可以是一個表格或一個函數。對於具有連續（或接近連續，例如像素）狀態的離散動作空間，習慣上使用非線性模型（例如神經網路）進行映射。Q-Value 網路的語義 (semantic) 希望非常簡單：我們只需要提供一個 tensor-to-tensor 映射，給定一個特定的狀態（輸入 tensor），輸出一個動作價值列表以供選擇。封裝器 (wrapper) 會將產生的動作連同動作價值列表寫入輸入的 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)

Distributional Q-learning 略有不同：在這種情況下，價值網路 (value network) 不會為每個狀態-動作價值輸出一個純量值。相反地，價值空間被劃分為任意數量的「區間 (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-Value actor 類別。
`QValueModule`(args, *kwargs)	用於 Q-value policies 的 Q-Value TensorDictModule。
`DistributionalQValueActor`(args, *kwargs)	一個 Distributional DQN actor 類別。
`DistributionalQValueModule`(args, *kwargs)	用於 Q-value policies 的 Distributional Q-Value hook。

價值運算子和聯合模型¶

TorchRL 提供了一系列的價值運算子 (value operators)，這些運算子封裝了價值網路，以簡化與程式庫其餘部分的介面。基本的構建模組是 torchrl.modules.tensordict_module.ValueOperator：給定一個輸入狀態（以及可能的動作），它會自動將 "state_value"（或 "state_action_value"）寫入 tensordict 中，具體取決於輸入是什麼。因此，這個類別同時考慮了價值網路和品質網路 (quality networks)。還提出了三個類別，將 policy 和價值網路組合在一起。ActorCriticOperator 是一個具有共享參數的聯合 actor-quality 網路：它讀取一個觀察，將其傳遞到一個共同的 backbone，寫入一個隱藏狀態，將這個隱藏狀態饋送到 policy，然後取得隱藏狀態和動作並提供狀態-動作對的品質。ActorValueOperator 是一個具有共享參數的聯合 actor-value 網路：它讀取一個觀察，將其傳遞到一個共同的 backbone，寫入一個隱藏狀態，將這個隱藏狀態饋送到 policy 和 value 模組以輸出一個動作和一個狀態價值。最後，ActorCriticWrapper 是一個沒有共享參數的聯合 actor 和 value 網路。它主要用來替代 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, ...)	Decision Transformer 的推論動作封裝器。

特定領域的 TensorDict 模組¶

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

`LMHeadActorValueOperator`(args, *kwargs)	從 huggingface-like *LMHeadModel 建構一個 Actor-Value 運算子。
`WorldModelWrapper`(args, *kwargs)	世界模型封裝器。

Hooks¶

Q-value hooks 被 QValueActor 和 DistributionalQValueActor 模組使用，通常應該優先考慮它們，因為它們更容易建立和使用。

`QValueHook`(action_space[, var_nums, ...])	用於 Q 值策略的 Q 值 Hook。
`DistributionalQValueHook`(action_space, support)	用於 Q-value policies 的 Distributional Q-Value hook。

模型¶

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[, ...])	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[, ...])	線上 Decion Transformer。
`DistributionalDQNnet`(args, *kwargs)	Distributional Deep Q-Network softmax 層。
`DreamerActor`(out_features[, depth, ...])	Dreamer actor 網路。
`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 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, ...])	實現具有位置縮放的常態分配。
`NormalParamWrapper`(operator[, ...])	常態分配參數的封裝器。
`TanhNormal`(loc, scale[, upscale, low, high, ...])	實現具有位置縮放的 TanhNormal 分配。
`TruncatedNormal`(loc, scale[, upscale, low, ...])	實現具有位置縮放的截斷常態分配。
`TanhDelta`(param[, low, high, event_dims, ...])	實現 Tanh transformed_in 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。