GRUModule¶

class torchrl.modules.GRUModule(*args, **kwargs)[source]¶

GRU 模块的嵌入器。

此类为 torch.nn.GRU 添加了以下功能：

与 TensorDict 的兼容性：隐藏状态被重塑以匹配 tensordict 的批次大小。
可选的多步执行：使用 torch.nn，必须在 torch.nn.GRUCell 和 torch.nn.GRU 之间进行选择，前者与单步输入兼容，后者与多步兼容。此类支持这两种用法。

构造完成后，模块**不会**设置为循环模式，即它将期望单步输入。

如果处于循环模式，则期望 tensordict 的最后一个维度标记步数。tensordict 的维度没有任何限制（除了对于时间输入，它必须大于一）。

参数:

input_size – 输入 x 中预期特征的数量
hidden_size – 隐藏状态 h 中的特征数量
num_layers – 循环层数。例如，设置 num_layers=2 将意味着堆叠两个 GRU 以形成一个堆叠 GRU，第二个 GRU 接收第一个 GRU 的输出并计算最终结果。默认值：1
bias – 如果为 False，则该层不使用偏置权重。默认值：True
dropout – 如果非零，则在除最后一层外的每个 GRU 层的输出上引入Dropout 层，dropout 概率等于dropout。默认值：0
python_based – 如果为 True，将使用 GRU 单元的完整 Python 实现。默认值：False

关键字参数:

in_key (str 或 tuple of str) – 模块的输入键。与 in_keys 互斥使用。如果提供，循环键假定为 [“recurrent_state”]，并且 in_key 将在此之前附加。
in_keys (list of str) – 一对字符串，对应于输入值和循环条目。与 in_key 互斥。
out_key (str 或 tuple of str) – 模块的输出键。与 out_keys 互斥使用。如果提供，循环键假定为 [(“recurrent_state”)]，并且 out_key 将在此之前附加。

out_keys (list of str) –

一对字符串，对应于输出值、第一个和第二个隐藏键。 .. note

For a better integration with TorchRL's environments, the best naming
for the output hidden key is ``("next", <custom_key>)``, such
that the hidden values are passed from step to step during a rollout.

device (torch.device 或 compatible) – 模块的设备。
gru (torch.nn.GRU, optional) – 要包装的 GRU 实例。与其他 nn.GRU 参数互斥。
default_recurrent_mode (bool, optional) – 如果提供，则为循环模式，如果尚未被 set_recurrent_mode 上下文管理器/装饰器覆盖。默认为 False。

变量:

recurrent_mode – 返回模块的循环模式。

set_recurrent_mode()[source]¶: 控制模块是否应以循环模式执行。

make_tensordict_primer()[source]¶: 为环境创建 TensorDictPrimer 变换，以使其感知 RNN 的循环状态。

注意

此模块依赖于输入 TensorDicts 中存在的特定 recurrent_state 键。要生成一个 TensorDictPrimer 变换，它将自动向环境 TensorDicts 添加隐藏状态，请使用方法 make_tensordict_primer()。如果此类是大型模块中的子模块，则可以通过父模块调用方法 get_primers_from_module() 来自动生成此模块（包括此模块）所需的所有 primer 变换。

示例

>>> from torchrl.envs import TransformedEnv, InitTracker
>>> from torchrl.envs import GymEnv
>>> from torchrl.modules import MLP
>>> from torch import nn
>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), InitTracker())
>>> gru_module = GRUModule(
...     input_size=env.observation_spec["observation"].shape[-1],
...     hidden_size=64,
...     in_keys=["observation", "rs"],
...     out_keys=["intermediate", ("next", "rs")])
>>> mlp = MLP(num_cells=[64], out_features=1)
>>> policy = Seq(gru_module, Mod(mlp, in_keys=["intermediate"], out_keys=["action"]))
>>> policy(env.reset())
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        intermediate: Tensor(shape=torch.Size([64]), device=cpu, dtype=torch.float32, is_shared=False),
        is_init: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                rs: Tensor(shape=torch.Size([1, 64]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)
>>> gru_module_training = gru_module.set_recurrent_mode()
>>> policy_training = Seq(gru_module, Mod(mlp, in_keys=["intermediate"], out_keys=["action"]))
>>> traj_td = env.rollout(3) # some random temporal data
>>> traj_td = policy_training(traj_td)
>>> print(traj_td)
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        intermediate: Tensor(shape=torch.Size([3, 64]), device=cpu, dtype=torch.float32, is_shared=False),
        is_init: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                is_init: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
                rs: Tensor(shape=torch.Size([3, 1, 64]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([3]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([3]),
    device=cpu,
    is_shared=False)

forward(tensordict: TensorDictBase = None)[source]¶

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

所有子类都应重写此方法。

注意

尽管前向传播的实现需要在该函数内定义，但之后应调用 Module 实例而不是此函数，因为前者负责运行已注册的钩子，而后者会默默忽略它们。

make_cudnn_based() → GRUModule[source]¶

将 GRU 层转换为其 CuDNN 版本。

返回:: self

make_python_based() → GRUModule[source]¶

将 GRU 层转换为其基于 Python 的版本。

返回:: self

make_tensordict_primer()[source]¶

为环境创建一个 tensordict primer。

一个 TensorDictPrimer 对象将确保策略在 rolloutexecution 过程中了解附加输入和输出（循环状态）。这样，数据就可以跨进程共享并得到妥善处理。

不在环境中包含 TensorDictPrimer 可能会导致行为定义不当，例如在并行设置中，一个步骤涉及将新的循环状态从 "next" 复制到根 tensordict，而 meth:~torchrl.EnvBase.step_mdp 方法将无法执行此操作，因为循环状态未在环境规范中注册。

在使用批处理环境（如 ParallelEnv）时，该变换可以在单个环境实例级别（即，一组具有内部设置的 tensordict primer 的转换后的环境）或在批处理环境实例级别（即，一组转换后的常规环境）使用。

有关为给定模块生成所有 primer 的方法，请参阅 torchrl.modules.utils.get_primers_from_module()。

示例

>>> from torchrl.collectors import SyncDataCollector
>>> from torchrl.envs import TransformedEnv, InitTracker
>>> from torchrl.envs import GymEnv
>>> from torchrl.modules import MLP, LSTMModule
>>> from torch import nn
>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
>>>
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), InitTracker())
>>> gru_module = GRUModule(
...     input_size=env.observation_spec["observation"].shape[-1],
...     hidden_size=64,
...     in_keys=["observation", "rs"],
...     out_keys=["intermediate", ("next", "rs")])
>>> mlp = MLP(num_cells=[64], out_features=1)
>>> policy = Seq(gru_module, Mod(mlp, in_keys=["intermediate"], out_keys=["action"]))
>>> policy(env.reset())
>>> env = env.append_transform(gru_module.make_tensordict_primer())
>>> data_collector = SyncDataCollector(
...     env,
...     policy,
...     frames_per_batch=10
... )
>>> for data in data_collector:
...     print(data)
...     break

GRUModule¶

文档

教程

资源