Recurrent Policies with LSTM and GRU Modules in TorchRL

import torch
import torch.nn as nn
from tensordict.nn import TensorDictModule, TensorDictSequential
from torchrl.envs import GymEnv, TransformedEnv, InitTracker
from torchrl.modules import LSTMModule, MLP

# Build the base environment with InitTracker
# InitTracker adds an "is_init" flag that tells LSTMModule when to zero hidden states
base_env = GymEnv("Pendulum-v1")
env = TransformedEnv(base_env, InitTracker())

obs_dim = env.observation_spec["observation"].shape[-1]   # 3 for Pendulum-v1
hidden_size = 64
action_dim = env.action_spec.shape[-1]                    # 1 for Pendulum-v1

# LSTMModule: obs → hidden representation + carries (h, c) across steps
lstm_module = LSTMModule(
    input_size=obs_dim,
    hidden_size=hidden_size,
    num_layers=1,
    # in_keys: [input_feature, hidden_h, hidden_c]
    in_keys=["observation", "rs_h", "rs_c"],
    # out_keys: [output_feature, next_hidden_h, next_hidden_c]
    out_keys=["lstm_out", ("next", "rs_h"), ("next", "rs_c")],
)

# A simple head MLP that maps the LSTM output to an action
mlp_head = TensorDictModule(
    MLP(in_features=hidden_size, out_features=action_dim, num_cells=[64]),
    in_keys=["lstm_out"],
    out_keys=["action"],
)

# Chain LSTM and MLP into a single policy
policy = TensorDictSequential(lstm_module, mlp_head)

from torchrl.envs import TransformedEnv

# make_tensordict_primer reads the LSTMModule's out_keys to determine
# what shape and dtype to pre-allocate
primer = lstm_module.make_tensordict_primer()
env = env.append_transform(primer)

# Alternatively, for a policy composed of multiple modules:
from torchrl.modules.utils import get_primers_from_module
primers = get_primers_from_module(policy)
for p in primers:
    env = env.append_transform(p)

# Verify everything is wired up
from torchrl.envs.utils import check_env_specs
check_env_specs(env)

from torchrl.modules import GRUModule

gru_module = GRUModule(
    input_size=obs_dim,
    hidden_size=hidden_size,
    num_layers=1,
    # in_keys: [input_feature, hidden_state]
    in_keys=["observation", "rs"],
    # out_keys: [output_feature, next_hidden_state]
    out_keys=["gru_out", ("next", "rs")],
)

mlp_head_gru = TensorDictModule(
    MLP(in_features=hidden_size, out_features=action_dim, num_cells=[64]),
    in_keys=["gru_out"],
    out_keys=["action"],
)

gru_policy = TensorDictSequential(gru_module, mlp_head_gru)

# Add primer for the GRU hidden state
env_gru = TransformedEnv(GymEnv("Pendulum-v1"), InitTracker())
env_gru = env_gru.append_transform(gru_module.make_tensordict_primer())

from torchrl.collectors import Collector

FRAMES_PER_BATCH = 1024
TOTAL_FRAMES = 100_000

collector = Collector(
    create_env_fn=lambda: TransformedEnv(
        GymEnv("Pendulum-v1"),
        InitTracker(),
    ).append_transform(lstm_module.make_tensordict_primer()),
    policy=policy,
    frames_per_batch=FRAMES_PER_BATCH,
    total_frames=TOTAL_FRAMES,
    device="cpu",
)

# The collector automatically handles:
# - resetting hidden states at episode boundaries (via "is_init")
# - copying ("next", "rs_h") → "rs_h" between steps
for batch in collector:
    # batch["rs_h"] and batch["rs_c"] are the per-step hidden states
    # batch[("next", "rs_h")] is the hidden state after each step
    print(batch.shape, batch.keys())
    break

collector.shutdown()

from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
from torchrl.objectives import ClipPPOLoss
from torchrl.objectives.value.advantages import GAE
from torchrl.modules import ValueOperator, set_recurrent_mode

# Critic that also takes hidden state as input
# (simplest option: feedforward critic on LSTM output)
lstm_critic = LSTMModule(
    input_size=obs_dim,
    hidden_size=hidden_size,
    in_keys=["observation", "critic_rs_h", "critic_rs_c"],
    out_keys=["critic_lstm_out", ("next", "critic_rs_h"), ("next", "critic_rs_c")],
)
value_head = TensorDictModule(
    MLP(in_features=hidden_size, out_features=1, num_cells=[64]),
    in_keys=["critic_lstm_out"],
    out_keys=["state_value"],
)
critic = TensorDictSequential(lstm_critic, value_head)

adv_module = GAE(gamma=0.99, lmbda=0.95, value_network=critic)
loss_module = ClipPPOLoss(
    actor_network=policy,
    critic_network=critic,
    clip_epsilon=0.2,
    entropy_coeff=0.01,
)

replay_buffer = TensorDictReplayBuffer(
    storage=LazyTensorStorage(FRAMES_PER_BATCH),
    sampler=SamplerWithoutReplacement(),
    batch_size=128,
)
optim = torch.optim.Adam(loss_module.parameters(), lr=3e-4)

for batch in collector:
    # Switch LSTM to recurrent mode for advantage computation
    with set_recurrent_mode(True), torch.no_grad():
        batch = adv_module(batch)

    replay_buffer.extend(batch.reshape(-1))

    for minibatch in replay_buffer:
        optim.zero_grad(set_to_none=True)
        loss = loss_module(minibatch)
        (loss["loss_objective"] + loss["loss_critic"] + loss["loss_entropy"]).backward()
        torch.nn.utils.clip_grad_norm_(loss_module.parameters(), 0.5)
        optim.step()

    collector.update_policy_weights_()

# Compile-friendly setup with the scan backend
lstm_scan = LSTMModule(
    input_size=obs_dim,
    hidden_size=hidden_size,
    in_keys=["observation", "rs_h", "rs_c"],
    out_keys=["lstm_out", ("next", "rs_h"), ("next", "rs_c")],
    recurrent_backend="scan",     # requires: pip install hoptorch
)

# CUDA Triton backend with gradient checkpointing
lstm_triton = LSTMModule(
    input_size=obs_dim,
    hidden_size=hidden_size,
    in_keys=["observation", "rs_h", "rs_c"],
    out_keys=["lstm_out", ("next", "rs_h"), ("next", "rs_c")],
    recurrent_backend="triton",   # requires: CUDA + pip install triton>=2.2
    recurrent_recompute="full",   # drop gate buffers from autograd graph
    recurrent_matmul_precision="tf32",  # fastest on Ampere+
)