ProbabilisticTensorDictModule

class tensordict.nn.ProbabilisticTensorDictModule(*args, **kwargs)

一个概率性 TD 模块。

ProbabilisticTensorDictModule 是一个非参数模块，嵌入了一个概率分布构造器。它使用指定的 in_keys 从输入的 TensorDict 中读取分布参数，并输出该分布的一个样本（非严格意义上）。

输出的“样本”是根据特定规则生成的，该规则由输入的 default_interaction_type 参数和 interaction_type() 全局函数指定。

ProbabilisticTensorDictModule 可用于构建分布（通过 get_dist() 方法）和/或从该分布中进行采样（通过对模块进行常规的 __call__() 调用）。

一个 ProbabilisticTensorDictModule 实例具有两个主要特性

• 它可以从 TensorDict 对象读取和写入数据；

• 它使用一个实值映射 R^n -> R^m 来在 R^d 中创建一个分布，可以从中采样或计算值。

当调用 __call__() 和 forward() 方法时，会创建一个分布并计算一个值（取决于 interaction_type 的值，可以使用 'dist.mean'、'dist.mode'、'dist.median' 属性，以及 'dist.rsample'、'dist.sample' 方法）。如果提供的 TensorDict 已经包含所有期望的键值对，则会跳过采样步骤。

默认情况下，ProbabilisticTensorDictModule 的分布类是 Delta 分布，这使得 ProbabilisticTensorDictModule 成为确定性映射函数的一个简单包装器。

参数：

• in_keys (NestedKey | List[NestedKey] | Dict[str, NestedKey]) – 将从输入的 TensorDict 中读取并用于构建分布的键。重要的是，如果它是 NestedKey 列表或单个 NestedKey，这些键的叶子（最后一个元素）必须与感兴趣的分布类使用的关键字匹配，例如 "loc" 和 "scale" 对于 Normal 分布等。如果 in_keys 是一个字典，则字典的键是分布的键，值是 tensordict 中将与相应分布键匹配的键。

• out_keys (NestedKey | List[NestedKey] | None) – 将写入采样值的键。重要的是，如果在输入的 TensorDict 中找到了这些键，则会跳过采样步骤。

关键字参数：

• default_interaction_type (InteractionType, optional) –

  仅限关键字参数。用于检索输出值的默认方法。应为 InteractionType 中的一个：MODE、MEDIAN、MEAN 或 RANDOM（在这种情况下，值从分布中随机采样）。默认值是 MODE。

  注意

  当抽取样本时，ProbabilisticTensorDictModule 实例将首先查找由 interaction_type() 全局函数指定的交互模式。如果此函数返回 None（其默认值），则将使用 ProbabilisticTDModule 实例的 default_interaction_type。请注意，DataCollectorBase 实例默认将 set_interaction_type 设置为 tensordict.nn.InteractionType.RANDOM。

  注意

  在某些情况下，模式、中位数或均值可能无法通过相应的属性直接获得。为解决此问题，ProbabilisticTensorDictModule 会首先尝试通过调用 get_mode()、get_median() 或 get_mean()（如果方法存在）来获取值。

• distribution_class (Type or Callable[[Any], Distribution], optional) –

  仅限关键字参数。用于采样的 torch.distributions.Distribution 类。默认值是 Delta。

  注意

  如果分布类是 CompositeDistribution 类型，则可以直接从此类的 distribution_kwargs 关键字参数中提供的 "distribution_map" 或 "name_map" 关键字参数推断出 out_keys，从而在这些情况下 out_keys 是可选的。

• distribution_kwargs (dict, optional) –

  仅限关键字参数。要传递给分布的关键字参数对。

  注意

  如果您的 kwargs 包含希望随模块一起传输到设备的张量，或者在调用 module.to(dtype) 时应修改其 dtype 的张量，您可以将 kwargs 包装在 TensorDictParams 中以自动完成此操作。

• return_log_prob (bool, optional) – 仅限关键字参数。如果为 True，则分布样本的对数概率将写入 tensordict 中，使用键 log_prob_key。默认值为 False。

• log_prob_keys (List[NestedKey], optional) –

  如果 return_log_prob=True，则写入 log_prob 的键。默认为 ‘<sample_key_name>_log_prob’，其中 <sample_key_name> 是 out_keys 中的每一个。

  注意

  这仅在 composite_lp_aggregate() 设置为 False 时可用。

• log_prob_key (NestedKey, optional) –

  如果 return_log_prob=True，则写入 log_prob 的键。当 composite_lp_aggregate() 设置为 True 时默认为 ‘sample_log_prob’，否则默认为 ‘<sample_key_name>_log_prob’。

  注意

  当有多个样本时，这仅在 composite_lp_aggregate() 设置为 True 时可用。

• cache_dist (bool, optional) – 仅限关键字参数。实验性：如果为 True，则分布的参数（即模块的输出）将与样本一起写入 tensordict。这些参数可用于稍后重新计算原始分布（例如，计算用于采样动作的分布与 PPO 中更新的分布之间的散度）。默认值为 False。

• n_empirical_estimate (int, optional) – 仅限关键字参数。当经验均值不可用时，用于计算经验均值的样本数量。默认为 1000。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import (
...     ProbabilisticTensorDictModule,
...     ProbabilisticTensorDictSequential,
...     TensorDictModule,
... )
>>> from tensordict.nn.distributions import NormalParamExtractor
>>> from tensordict.nn.functional_modules import make_functional
>>> from torch.distributions import Normal, Independent
>>> td = TensorDict(
...     {"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3]
... )
>>> net = torch.nn.GRUCell(4, 8)
>>> module = TensorDictModule(
...     net, in_keys=["input", "hidden"], out_keys=["params"]
... )
>>> normal_params = TensorDictModule(
...     NormalParamExtractor(), in_keys=["params"], out_keys=["loc", "scale"]
... )
>>> def IndepNormal(**kwargs):
...     return Independent(Normal(**kwargs), 1)
>>> prob_module = ProbabilisticTensorDictModule(
...     in_keys=["loc", "scale"],
...     out_keys=["action"],
...     distribution_class=IndepNormal,
...     return_log_prob=True,
... )
>>> td_module = ProbabilisticTensorDictSequential(
...     module, normal_params, prob_module
... )
>>> params = TensorDict.from_module(td_module)
>>> with params.to_module(td_module):
...     _ = td_module(td)
>>> print(td)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        hidden: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        input: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        loc: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        params: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> with params.to_module(td_module):
...     dist = td_module.get_dist(td)
>>> print(dist)
Independent(Normal(loc: torch.Size([3, 4]), scale: torch.Size([3, 4])), 1)
>>> # we can also apply the module to the TensorDict with vmap
>>> from torch import vmap
>>> params = params.expand(4)
>>> def func(td, params):
...     with params.to_module(td_module):
...         return td_module(td)
>>> td_vmap = vmap(func, (None, 0))(td, params)
>>> print(td_vmap)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        hidden: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        input: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        loc: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        params: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False),
        sample_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        scale: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=None,
    is_shared=False)

build_dist_from_params(tensordict: TensorDictBase) → Distribution

使用输入的 tensordict 中提供的参数创建一个 torch.distribution.Distribution 实例。

参数：

tensordict (TensorDictBase) – 包含分布参数的输入 tensordict。

返回：

使用输入的 tensordict 创建的 torch.distribution.Distribution 实例。

抛出：

TypeError – 如果输入的 tensordict 与分布关键字不匹配。

property dist_params_keys: List[NestedKey]

返回指向分布参数的所有键。

property dist_sample_keys: List[NestedKey]

返回指向分布样本的所有键。

forward(tensordict: TensorDictBase = None, tensordict_out: tensordict.base.TensorDictBase | None = None, _requires_sample: bool = True) → TensorDictBase

定义每次调用时执行的计算。

应被所有子类覆盖。

注意

虽然 forward pass 的实现需要在函数内部定义，但之后应该调用 Module 实例而不是此函数本身，因为前者负责运行注册的钩子，而后者会默默地忽略它们。

get_dist(tensordict: TensorDictBase) → Distribution

使用输入的 tensordict 中提供的参数创建一个 torch.distribution.Distribution 实例。

参数：

tensordict (TensorDictBase) – 包含分布参数的输入 tensordict。

返回：

使用输入的 tensordict 创建的 torch.distribution.Distribution 实例。

抛出：

TypeError – 如果输入的 tensordict 与分布关键字不匹配。

log_prob(tensordict, *, dist: Optional[Distribution] = None)

计算分布样本的对数概率。

参数：

• tensordict (TensorDictBase) – 包含分布参数的输入 tensordict。

• dist (torch.distributions.Distribution, optional) – 分布实例。默认为 None。如果为 None，则将使用 get_dist 方法计算分布。

返回：

表示分布样本对数概率的张量。