SafeProbabilisticModule

class torchrl.modules.tensordict_module.SafeProbabilisticModule(*args, **kwargs)

tensordict.nn.ProbabilisticTensorDictModule 的子类，它接受一个 TensorSpec 作为参数来控制输出域。

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

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

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

SafeProbabilisticModule 实例具有两个主要特性：

• 它读写 TensorDict 对象；

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

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

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

该类与 tensordict.nn.ProbabilisticTensorDictModule 的区别在于，它接受一个 spec 关键字参数，该参数可用于控制样本是否属于该分布。 safe 关键字参数控制是否应根据 spec 检查样本值。

参数:

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

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

• spec (TensorSpec) – 第一个输出张量的 spec。在调用 td_module.random() 以在目标空间生成随机值时使用。

关键字参数:

• safe (bool, optional) – 如果为 True，则样本的值将根据输入 spec 进行检查。由于探索策略或数值下溢/上溢问题，可能会发生超出域的采样。与 spec 参数一样，此检查仅针对分布样本进行，而不针对输入模块返回的其他张量。如果样本超出范围，它将使用 TensorSpec.project 方法投影回所需空间。默认为 False。

• 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 或 Callable[[Any], Distribution], optional) –

  仅关键字参数。一个 torch.distributions.Distribution 类，用于采样。默认为 Delta。

  注意

  如果分布类是 CompositeDistribution 类型，则 out_keys 可以直接从该类的 distribution_kwargs 关键字参数中提供的 "distribution_map" 或 "name_map" 推断出来，在这种情况下 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。

警告

运行检查需要时间！使用 safe=True 将保证样本在 spec 边界内，这取决于 project() 中编码的一些启发式方法，但这就需要检查值是否在 spec 空间内，这会带来一些开销。

另请参阅

tensordict.nn.CompositeDistribution 可用于创建多头策略。

示例

>>> from torchrl.modules import SafeProbabilisticModule
>>> from torchrl.data import Bounded
>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import InteractionType
>>> mod = SafeProbabilisticModule(
...     in_keys=["loc", "scale"],
...     out_keys=["action"],
...     distribution_class=torch.distributions.Normal,
...     safe=True,
...     spec=Bounded(low=-1, high=1, shape=()),
...     default_interaction_type=InteractionType.RANDOM
... )
>>> _ = torch.manual_seed(0)
>>> data = TensorDict(
...     loc=torch.zeros(10, requires_grad=True),
...     scale=torch.full((10,), 10.0),
...     batch_size=(10,))
>>> data = mod(data)
>>> print(data["action"]) # All actions are within bound
tensor([ 1., -1., -1.,  1., -1., -1.,  1.,  1., -1., -1.],
       grad_fn=<ClampBackward0>)
>>> data["action"].mean().backward()
>>> print(data["loc"].grad) # clamp anihilates gradients
tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])

random(tensordict: TensorDictBase) → TensorDictBase

在目标空间中采样一个随机元素，而不考虑任何输入。

如果存在多个输出键，则只有第一个键将被写入输入的 tensordict 中。

参数:

tensordict (TensorDictBase) – 应将输出值写入的 tensordict。

返回:

带有输出键的新/更新值的原始 tensordict。

random_sample(tensordict: TensorDictBase) → TensorDictBase

请参阅 SafeModule.random(...)。