快捷方式

Actor

class torchrl.modules.tensordict_module.Actor(*args, **kwargs)[源码]

强化学习中用于确定性 Actor 的通用类。

Actor 类带有默认的 out_keys 值 (["action"]),并且如果提供了 spec 但不是 Composite 对象,它将自动转换为 spec = Composite(action=spec)

参数:
  • module (nn.Module) – 一个 Module,用于将输入映射到输出参数空间。

  • in_keys (字符串可迭代对象, 可选) – 从输入 tensordict 读取并传递给 module 的键。如果包含多个元素,值将按照 in_keys 迭代器给定的顺序传递。默认为 ["observation"]

  • out_keys (字符串可迭代对象) – 要写入输入 tensordict 的键。out_keys 的长度必须与嵌入的 module 返回的张量数量匹配。使用 "_" 作为键可以避免将张量写入输出。默认为 ["action"]

关键字参数:
  • spec (TensorSpec, 可选) – 仅限关键字参数。输出张量的规格。如果 module 输出多个张量,spec 描述第一个输出张量的空间。

  • safe (bool) – 仅限关键字参数。如果为 True,输出值将对照输入规格进行检查。由于探索策略或数值下溢/溢出问题,可能会发生域外采样。如果此值超出范围,将使用 project() 方法将其投影回期望空间。默认为 False

示例

>>> import torch
>>> from tensordict import TensorDict
>>> from torchrl.data import Unbounded
>>> from torchrl.modules import Actor
>>> torch.manual_seed(0)
>>> td = TensorDict({"observation": torch.randn(3, 4)}, [3,])
>>> action_spec = Unbounded(4)
>>> module = torch.nn.Linear(4, 4)
>>> td_module = Actor(
...    module=module,
...    spec=action_spec,
...    )
>>> td_module(td)
TensorDict(
    fields={
        action: 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)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> print(td.get("action"))
tensor([[-1.3635, -0.0340,  0.1476, -1.3911],
        [-0.1664,  0.5455,  0.2247, -0.4583],
        [-0.2916,  0.2160,  0.5337, -0.5193]], grad_fn=<AddmmBackward0>)

文档

查阅 PyTorch 全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源