LSTMModule¶

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

LSTM 模块的嵌入器。

此类为 torch.nn.LSTM 添加以下功能

与 TensorDict 兼容：隐藏状态会重新整形以匹配 TensorDict 的批次大小。
可选的多步执行：使用 torch.nn，必须在 torch.nn.LSTMCell 和 torch.nn.LSTM 之间进行选择，前者与单步输入兼容，后者与多步输入兼容。此类启用这两种用法。

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

如果处于循环模式，则预期 TensorDict 的最后一个维度表示步数。TensorDict 的维度没有约束（除了对于时间输入，它必须大于 1）。

注意

此类可以处理沿时间维度连续的多条轨迹，但在这些情况下不应信任最终的隐藏值（即，不应将它们重新用于连续的轨迹）。原因是 LSTM 只返回最后一个隐藏值，对于我们提供的填充输入，它可能对应于 0 填充的输入。

参数：

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

关键字参数：

in_key (str 或 str 元组) – 模块的输入键。与 in_keys 互斥。如果提供，则循环键假定为 [“recurrent_state_h”, “recurrent_state_c”]，并且 in_key 将附加到这些键之前。
in_keys (str 列表) – 对应于输入值、第一个和第二个隐藏键的三个字符串。与 in_key 互斥。
out_key (str 或 str 元组) – 模块的输出键。与 out_keys 互斥。如果提供，则循环键假定为 [(“next”, “recurrent_state_h”), (“next”, “recurrent_state_c”)]，并且 out_key 将附加到这些键之前。

out_keys (str 列表) –

对应于输出值、第一个和第二个隐藏键的三个字符串。.. 注意

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 或兼容) – 模块的设备。
lstm (torch.nn.LSTM, 可选) – 要包装的 LSTM 实例。与其他 nn.LSTM 参数互斥。

变量：

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

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

make_tensordict_primer()[source]¶: 为环境创建 TensorDictPrimer 转换，以了解 RNN 的循环状态。

注意

此模块依赖于输入 TensorDict 中存在特定的 recurrent_state 键。要生成一个 TensorDictPrimer 转换，该转换将自动将隐藏状态添加到环境 TensorDict 中，请使用方法 make_tensordict_primer()。如果此类是较大模块中的子模块，则可以在父模块上调用方法 get_primers_from_module() 以自动生成所有子模块（包括此模块）所需的引导转换。

示例

>>> 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())
>>> lstm_module = LSTMModule(
...     input_size=env.observation_spec["observation"].shape[-1],
...     hidden_size=64,
...     in_keys=["observation", "rs_h", "rs_c"],
...     out_keys=["intermediate", ("next", "rs_h"), ("next", "rs_c")])
>>> mlp = MLP(num_cells=[64], out_features=1)
>>> policy = Seq(lstm_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_c: Tensor(shape=torch.Size([1, 64]), device=cpu, dtype=torch.float32, is_shared=False),
                rs_h: 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)

forward(tensordict: TensorDictBase)[source]¶

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

应由所有子类覆盖。

注意

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

make_tensordict_primer()[source]¶

为环境创建 TensorDict 引导程序。

一个 TensorDictPrimer 对象将确保策略在回滚执行期间了解补充输入和输出（循环状态）。这样，数据可以在进程之间共享并得到正确处理。

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

示例

>>> 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())
>>> lstm_module = LSTMModule(
...     input_size=env.observation_spec["observation"].shape[-1],
...     hidden_size=64,
...     in_keys=["observation", "rs_h", "rs_c"],
...     out_keys=["intermediate", ("next", "rs_h"), ("next", "rs_c")])
>>> mlp = MLP(num_cells=[64], out_features=1)
>>> policy = Seq(lstm_module, Mod(mlp, in_keys=["intermediate"], out_keys=["action"]))
>>> policy(env.reset())
>>> env = env.append_transform(lstm_module.make_tensordict_primer())
>>> data_collector = SyncDataCollector(
...     env,
...     policy,
...     frames_per_batch=10
... )
>>> for data in data_collector:
...     print(data)
...     break

set_recurrent_mode(mode: bool = True)[source]¶

返回模块的新副本，该副本共享相同的 lstm 模型，但具有不同的 recurrent_mode 属性（如果不同）。

创建副本是为了使模块可以在代码的不同部分（推理与训练）中使用不同的行为。

示例

>>> from torchrl.envs import TransformedEnv, InitTracker, step_mdp
>>> from torchrl.envs import GymEnv
>>> from torchrl.modules import MLP
>>> from tensordict import TensorDict
>>> from torch import nn
>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
>>> env = TransformedEnv(GymEnv("Pendulum-v1"), InitTracker())
>>> lstm = nn.LSTM(input_size=env.observation_spec["observation"].shape[-1], hidden_size=64, batch_first=True)
>>> lstm_module = LSTMModule(lstm=lstm, in_keys=["observation", "hidden0", "hidden1"], out_keys=["intermediate", ("next", "hidden0"), ("next", "hidden1")])
>>> mlp = MLP(num_cells=[64], out_features=1)
>>> # building two policies with different behaviours:
>>> policy_inference = Seq(lstm_module, Mod(mlp, in_keys=["intermediate"], out_keys=["action"]))
>>> policy_training = Seq(lstm_module.set_recurrent_mode(True), Mod(mlp, in_keys=["intermediate"], out_keys=["action"]))
>>> traj_td = env.rollout(3) # some random temporal data
>>> traj_td = policy_training(traj_td)
>>> # let's check that both return the same results
>>> td_inf = TensorDict({}, traj_td.shape[:-1])
>>> for td in traj_td.unbind(-1):
...     td_inf = td_inf.update(td.select("is_init", "observation", ("next", "observation")))
...     td_inf = policy_inference(td_inf)
...     td_inf = step_mdp(td_inf)
...
>>> torch.testing.assert_close(td_inf["hidden0"], traj_td[..., -1]["next", "hidden0"])

LSTMModule¶

文档

教程

资源