ProbabilisticActor¶
- class torchrl.modules.tensordict_module.ProbabilisticActor(*args, **kwargs)[来源]¶
RL 中概率型 actor 的通用类。
Actor 类带有 out_keys ([“action”]) 的默认值,如果提供了 spec 但不是 Composite 对象,它将被自动转换为
spec = Composite(action=spec)
- 参数:
module (nn.Module) – 用于将输入映射到输出参数空间的
torch.nn.Module
实例。in_keys (str 或 str 的可迭代对象 或 dict) – 将从输入 TensorDict 中读取并用于构建分布的键。重要的是,如果它是一个字符串的可迭代对象或一个字符串,这些键必须与感兴趣的分布类使用的关键字匹配,例如 Normal 分布的
"loc"
和"scale"
等。如果 in_keys 是一个字典,键是分布的键,值是 tensordict 中将与对应分布键匹配的键。out_keys (str 或 str 的可迭代对象) – 采样值将被写入的键。重要的是,如果在输入 TensorDict 中找到这些键,采样步骤将被跳过。
spec (TensorSpec, 可选) – 仅关键字参数,包含输出张量的 spec。如果模块输出多个张量,spec 描述第一个输出张量的空间。
safe (bool) – 仅关键字参数。如果为
True
,则对照输入 spec 检查输出值。由于探索策略或数值下溢/上溢问题,可能会发生域外采样。如果此值超出范围,则使用TensorSpec.project
方法将其投影回期望的空间。默认为False
。default_interaction_type (tensordict.nn.InteractionType, 可选) –
仅关键字参数。用于检索输出值的默认方法。应为以下之一:
InteractionType.MODE
、InteractionType.DETERMINISTIC
、InteractionType.MEDIAN
、InteractionType.MEAN
或InteractionType.RANDOM
(在这种情况下,值会从分布中随机采样)。TorchRL 的ExplorationType
类是InteractionType
的代理。默认为InteractionType.DETERMINISTIC
。注意
当抽取样本时,
ProbabilisticActor
实例将首先查找由interaction_type()
全局函数指定的交互模式。如果此函数返回 None(其默认值),则将使用 ProbabilisticTDModule 实例的 default_interaction_type。请注意,DataCollectorBase
实例默认会将 set_interaction_type 设置为tensordict.nn.InteractionType.RANDOM
。distribution_class (Type, 可选) –
仅关键字参数。用于采样的
torch.distributions.Distribution
类。默认为tensordict.nn.distributions.Delta
。注意
如果
distribution_class
的类型为CompositeDistribution
,则键将从该分布的distribution_map
/name_map
关键字参数中推断出来。如果此分布与另一个构造函数(例如 partial 或 lambda 函数)一起使用,则需要显式提供 out_keys。另请注意,动作将不会带有"action"
键前缀,请参见下面的示例,了解如何使用ProbabilisticActor
实现此功能。distribution_kwargs (dict, 可选) – 仅关键字参数。将传递给分布的关键字参数对。
return_log_prob (bool, 可选) – 仅关键字参数。如果为
True
,则分布样本的对数概率将以键 ‘sample_log_prob’ 写入 tensordict 中。默认为False
。cache_dist (bool, 可选) – 仅关键字参数。实验性功能:如果为
True
,则分布的参数(即模块的输出)将与样本一起写入 tensordict。这些参数可用于以后重新计算原始分布(例如,计算用于采样动作的分布与 PPO 中更新的分布之间的散度)。默认为False
。n_empirical_estimate (int, 可选) – 仅关键字参数。当经验均值不可用时,用于计算的样本数。默认为 1000。
示例
>>> import torch >>> from tensordict import TensorDict >>> from tensordict.nn import TensorDictModule >>> from torchrl.data import Bounded >>> from torchrl.modules import ProbabilisticActor, NormalParamExtractor, TanhNormal >>> td = TensorDict({"observation": torch.randn(3, 4)}, [3,]) >>> action_spec = Bounded(shape=torch.Size([4]), ... low=-1, high=1) >>> module = nn.Sequential(torch.nn.Linear(4, 8), NormalParamExtractor()) >>> tensordict_module = TensorDictModule(module, in_keys=["observation"], out_keys=["loc", "scale"]) >>> td_module = ProbabilisticActor( ... module=tensordict_module, ... spec=action_spec, ... in_keys=["loc", "scale"], ... distribution_class=TanhNormal, ... ) >>> td = td_module(td) >>> td TensorDict( fields={ action: 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), observation: Tensor(shape=torch.Size([3, 4]), 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)
概率型 actor 也支持通过
tensordict.nn.CompositeDistribution
类实现复合动作。此分布以 tensordict 作为输入(通常为 “params”)并将其作为一个整体读取:此 tensordict 的内容是复合分布中包含的各个分布的输入。示例
>>> from tensordict import TensorDict >>> from tensordict.nn import CompositeDistribution, TensorDictModule >>> from torchrl.modules import ProbabilisticActor >>> from torch import nn, distributions as d >>> import torch >>> >>> class Module(nn.Module): ... def forward(self, x): ... return x[..., :3], x[..., 3:6], x[..., 6:] >>> module = TensorDictModule(Module(), ... in_keys=["x"], ... out_keys=[("params", "normal", "loc"), ... ("params", "normal", "scale"), ... ("params", "categ", "logits")]) >>> actor = ProbabilisticActor(module, ... in_keys=["params"], ... distribution_class=CompositeDistribution, ... distribution_kwargs={"distribution_map": { ... "normal": d.Normal, "categ": d.Categorical}} ... ) >>> data = TensorDict({"x": torch.rand(10)}, []) >>> actor(data) TensorDict( fields={ categ: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), normal: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), params: TensorDict( fields={ categ: TensorDict( fields={ logits: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), normal: TensorDict( fields={ loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), scale: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), x: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)
使用复合分布的概率型 actor 可以通过以下示例代码实现
示例
>>> import torch >>> from tensordict import TensorDict >>> from tensordict.nn import CompositeDistribution >>> from tensordict.nn import TensorDictModule >>> from torch import distributions as d >>> from torch import nn >>> >>> from torchrl.modules import ProbabilisticActor >>> >>> >>> class Module(nn.Module): ... def forward(self, x): ... return x[..., :3], x[..., 3:6], x[..., 6:] ... >>> >>> module = TensorDictModule(Module(), ... in_keys=["x"], ... out_keys=[ ... ("params", "normal", "loc"), ("params", "normal", "scale"), ("params", "categ", "logits") ... ]) >>> actor = ProbabilisticActor(module, ... in_keys=["params"], ... distribution_class=CompositeDistribution, ... distribution_kwargs={"distribution_map": {"normal": d.Normal, "categ": d.Categorical}, ... "name_map": {"normal": ("action", "normal"), ... "categ": ("action", "categ")}} ... ) >>> print(actor.out_keys) [('params', 'normal', 'loc'), ('params', 'normal', 'scale'), ('params', 'categ', 'logits'), ('action', 'normal'), ('action', 'categ')] >>> >>> data = TensorDict({"x": torch.rand(10)}, []) >>> module(data) >>> print(actor(data)) TensorDict( fields={ action: TensorDict( fields={ categ: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False), normal: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), params: TensorDict( fields={ categ: TensorDict( fields={ logits: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), normal: TensorDict( fields={ loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), scale: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False), x: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([]), device=None, is_shared=False)