概率 TensorDict 模块

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 方法时，会创建一个分布并计算一个值（使用 'mean'、'mode'、'median' 属性或 'rsample'、'sample' 方法）。如果提供的 TensorDict 已经包含所有所需键值对，则跳过采样步骤。

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

参数：

• in_keys (NestedKey 或 NestedKey 列表 或 dict) – 将从输入 TensorDict 中读取并用于构建分布的键。重要的是，如果它是一个 NestedKey 列表或一个 NestedKey，则这些键的叶（最后一个元素）必须与目标分布类使用的关键字匹配，例如 "loc" 和 "scale" 用于正态分布等等。如果 in_keys 是一个字典，则键是分布的键，值是 tensordict 中将与相应分布键匹配的键。

• out_keys (NestedKey 或 NestedKey 列表) – 将写入采样值的键。重要的是，如果在输入 TensorDict 中找到这些键，则将跳过采样步骤。

• default_interaction_mode (str, 可选) – 已弃用 的关键字参数。请改用 default_interaction_type。

• default_interaction_type (InteractionType, 可选) –

  关键字参数。用于检索输出值的默认方法。应该是 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, 可选) –

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

  注意

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

• distribution_kwargs (dict, 可选) – 关键字参数。要传递给分布的关键字参数对。

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

• log_prob_key (NestedKey, 可选) – 如果 return_log_prob = True，则写入 log_prob 的键。默认为 ‘sample_log_prob’。

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

• n_empirical_estimate (int, 可选) – 关键字参数。当不可用时，用于计算经验平均值的样本数量。默认为 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)

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

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

应由所有子类覆盖。

注意

虽然前向传递的配方需要在此函数内定义，但应该在之后调用 Module 实例而不是此函数，因为前者负责运行已注册的钩子，而后者则静默地忽略它们。

get_dist(tensordict: TensorDictBase) → Distribution

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

log_prob(tensordict)

写入分布样本的对数概率。