Repository: pytorch/rl

Issue:
[BUG] LSTMModule sets hidden state to 0.0 at every step when combined with ParallelEnv.
## Describe the bug

While trying to use an actor network that contains an LSTMModule together with a vectorised environment, I detected the following:
 - When the vectorised environment is created using the SerialEnv class everything works fine.
 - When the vectorised environment is created using the ParallelEnv, the LSTM hidden states are not passed forward from step to step and are always 0.0. 
I suspect somewhere in the ParallelEnv the hidden states are being reset at every step.

## To Reproduce

Here is the complete code to reproduce the issue. Note that if in the method create_env_fn(), ParallelEnv is replaced by SerialEnv, the code works.

```python

import torch
import tqdm
from tensordict.nn import TensorDictModule, TensorDictSequential
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import ProbabilisticActor, LSTMModule, MLP
from torchrl.collectors import SyncDataCollector
from torchrl.envs import (
    ParallelEnv,
    SerialEnv,
    TransformedEnv,
    InitTracker,
    StepCounter,
    RewardSum,
)


def create_model(input_size, output_size, hidden_size=256, num_layers=3, out_key="logits"):

    lstm_module = LSTMModule(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        in_key="observation",
        out_key="features",
    )
    mlp = TensorDictModule(
        MLP(
            in_features=hidden_size,
            out_features=output_size,
            num_cells=[],
        ),
        in_keys=["features"],
        out_keys=[out_key],
    )

    return TensorDictSequential(lstm_module, mlp)


def create_rhs_transform(input_size, hidden_size=256, num_layers=3):
    lstm_module = LSTMModule(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        in_key="observation",
        out_key="features",
    )
    return lstm_module.make_tensordict_primer()


def main():

    device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
    test_env = GymEnv("CartPole-v1", device=device, categorical_action_encoding=True)
    action_spec = test_env.action_spec.space
    observation_spec = test_env.observation_spec["observation"]

    def create_transformed_env():
        env = GymEnv("CartPole-v1", device=device, categorical_action_encoding=True)
        env = TransformedEnv(env)
        env.append_transform(create_rhs_transform(input_size=observation_spec.shape[-1]))
        env.append_transform(StepCounter())
        env.append_transform(RewardSum())
        env.append_transform(InitTracker())
        return env

    def create_env_fn(num_workers=4):
        env = ParallelEnv(  # There is some bug here! When using ParallelEnv, hidden states are always zero
        # env = SerialEnv(  # This works!
            create_env_fn=create_transformed_env,
            num_workers=num_workers,
        )
        return env

    # Model
    ##################

    actor_model = create_model(input_size=observation_spec.shape[-1], output_size=action_spec.n)
    actor = ProbabilisticActor(
        module=actor_model,
        in_keys=["logits"],
        out_keys=["action"],
        distribution_class=torch.distributions.Categorical,
        return_log_prob=True,
    )
    actor = actor.to(device)

    # Collector
    ##################

    collector = SyncDataCollector(
            create_env_fn=create_env_fn,
            policy=actor,
            frames_per_batch=100,
            total_frames=1000,
            device=device,
            storing_device=device,
            max_frames_per_traj=-1,
            split_trajs=False,
        )

    # Collection loop
    ##################

    for data in tqdm.tqdm(collector):

        assert data["recurrent_state_c"].sum() != 0.0  # This is False when I use parallel envs
        assert data["next"]["recurrent_state_c"].sum() != 0.0  # This is False when I use parallel envs

    collector.shutdown()


if __name__ == "__main__":
    main()


```

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
Looking into this thanks!

Resolving patch:
diff --git a/torchrl/envs/common.py b/torchrl/envs/common.py
--- a/torchrl/envs/common.py
+++ b/torchrl/envs/common.py
@@ -1118,6 +1118,9 @@ def step(self, tensordict: TensorDictBase) -> TensorDictBase:
 
         Args:
             tensordict (TensorDictBase): Tensordict containing the action to be taken.
+                If the input tensordict contains a ``"next"`` entry, the values contained in it
+                will prevail over the newly computed values. This gives a mechanism
+                to override the underlying computations.
 
         Returns:
             the input tensordict, modified in place with the resulting observations, done state and reward
@@ -1126,10 +1129,13 @@ def step(self, tensordict: TensorDictBase) -> TensorDictBase:
         """
         # sanity check
         self._assert_tensordict_shape(tensordict)
+        next_preset = tensordict.get("next", None)
 
         next_tensordict = self._step(tensordict)
         next_tensordict = self._step_proc_data(next_tensordict)
-        # tensordict could already have a "next" key
+        if next_preset is not None:
+            # tensordict could already have a "next" key
+            next_tensordict.update(next_preset)
         tensordict.set("next", next_tensordict)
         return tensordict
 
@@ -1669,11 +1675,14 @@ def fake_tensordict(self) -> TensorDictBase:
         next_output.update(fake_reward)
         next_output.update(fake_done)
         fake_in_out.update(fake_done.clone())
+        if "next" not in fake_in_out.keys():
+            fake_in_out.set("next", next_output)
+        else:
+            fake_in_out.get("next").update(next_output)
 
-        fake_td = fake_in_out.set("next", next_output)
-        fake_td.batch_size = self.batch_size
-        fake_td = fake_td.to(self.device)
-        return fake_td
+        fake_in_out.batch_size = self.batch_size
+        fake_in_out = fake_in_out.to(self.device)
+        return fake_in_out
 
 
 class _EnvWrapper(EnvBase, metaclass=abc.ABCMeta):
diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py
--- a/torchrl/modules/tensordict_module/rnn.py
+++ b/torchrl/modules/tensordict_module/rnn.py
@@ -5,11 +5,11 @@
 from typing import Optional, Tuple
 
 import torch
-from tensordict import unravel_key_list
+from tensordict import TensorDictBase, unravel_key_list
 
 from tensordict.nn import TensorDictModuleBase as ModuleBase
 
-from tensordict.tensordict import NO_DEFAULT, TensorDictBase
+from tensordict.tensordict import NO_DEFAULT
 from tensordict.utils import prod
 
 from torch import nn


Test patch:
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
--- a/test/mocking_classes.py
+++ b/test/mocking_classes.py
@@ -421,12 +421,12 @@ def __new__(
                 shape=batch_size,
             )
         if action_spec is None:
-            action_spec_cls = (
-                DiscreteTensorSpec
-                if categorical_action_encoding
-                else OneHotDiscreteTensorSpec
-            )
-            action_spec = action_spec_cls(n=7, shape=(*batch_size, 7))
+            if categorical_action_encoding:
+                action_spec_cls = DiscreteTensorSpec
+                action_spec = action_spec_cls(n=7, shape=batch_size)
+            else:
+                action_spec_cls = OneHotDiscreteTensorSpec
+                action_spec = action_spec_cls(n=7, shape=(*batch_size, 7))
         if reward_spec is None:
             reward_spec = UnboundedContinuousTensorSpec(shape=(1,))
         if done_spec is None:
@@ -1053,7 +1053,7 @@ def _step(
             batch_size=self.batch_size,
             device=self.device,
         )
-        return tensordict.select().set("next", tensordict)
+        return tensordict
 
 
 class NestedCountingEnv(CountingEnv):
diff --git a/test/test_env.py b/test/test_env.py
--- a/test/test_env.py
+++ b/test/test_env.py
@@ -143,7 +143,7 @@ def test_env_seed(env_name, frame_skip, seed=0):
     env.set_seed(seed)
     td0b = env.fake_tensordict()
     td0b = env.reset(tensordict=td0b)
-    td1b = env.step(td0b.clone().set("action", action))
+    td1b = env.step(td0b.exclude("next").clone().set("action", action))
 
     assert_allclose_td(td0a, td0b.select(*td0a.keys()))
     assert_allclose_td(td1a, td1b)
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
--- a/test/test_tensordictmodules.py
+++ b/test/test_tensordictmodules.py
@@ -7,8 +7,14 @@
 
 import pytest
 import torch
+from mocking_classes import DiscreteActionVecMockEnv
 from tensordict import pad, TensorDict, unravel_key_list
-from tensordict.nn import InteractionType, make_functional, TensorDictModule
+from tensordict.nn import (
+    InteractionType,
+    make_functional,
+    TensorDictModule,
+    TensorDictSequential,
+)
 from torch import nn
 from torchrl.data.tensor_specs import (
     BoundedTensorSpec,
@@ -21,6 +27,7 @@
     DecisionTransformerInferenceWrapper,
     DTActor,
     LSTMModule,
+    MLP,
     NormalParamWrapper,
     OnlineDTActor,
     ProbabilisticActor,
@@ -1765,6 +1772,55 @@ def test_multi_consecutive(self, shape):
             td_ss["intermediate"], td["intermediate"][..., -1, :]
         )
 
+    def test_lstm_parallel_env(self):
+        from torchrl.envs import InitTracker, ParallelEnv, TransformedEnv
+
+        # tests that hidden states are carried over with parallel envs
+        lstm_module = LSTMModule(
+            input_size=7,
+            hidden_size=12,
+            num_layers=2,
+            in_key="observation",
+            out_key="features",
+        )
+
+        def create_transformed_env():
+            primer = lstm_module.make_tensordict_primer()
+            env = DiscreteActionVecMockEnv(categorical_action_encoding=True)
+            env = TransformedEnv(env)
+            env.append_transform(InitTracker())
+            env.append_transform(primer)
+            return env
+
+        env = ParallelEnv(
+            create_env_fn=create_transformed_env,
+            num_workers=2,
+        )
+
+        mlp = TensorDictModule(
+            MLP(
+                in_features=12,
+                out_features=7,
+                num_cells=[],
+            ),
+            in_keys=["features"],
+            out_keys=["logits"],
+        )
+
+        actor_model = TensorDictSequential(lstm_module, mlp)
+
+        actor = ProbabilisticActor(
+            module=actor_model,
+            in_keys=["logits"],
+            out_keys=["action"],
+            distribution_class=torch.distributions.Categorical,
+            return_log_prob=True,
+        )
+        for break_when_any_done in [False, True]:
+            data = env.rollout(10, actor, break_when_any_done=break_when_any_done)
+            assert (data.get("recurrent_state_c") != 0.0).any()
+            assert (data.get(("next", "recurrent_state_c")) != 0.0).all()
+
 
 def test_safe_specs():
 
diff --git a/test/test_transforms.py b/test/test_transforms.py
--- a/test/test_transforms.py
+++ b/test/test_transforms.py
@@ -430,10 +430,9 @@ def test_transform_env_clone(self):
         value_at_clone = td["next", "observation"].clone()
         for _ in range(10):
             td = env.rand_step(td)
-        assert (td["next", "observation"] != value_at_clone).any()
-        assert (
-            td["next", "observation"] == env.transform._cat_buffers_observation
-        ).all()
+            td = step_mdp(td)
+        assert (td["observation"] != value_at_clone).any()
+        assert (td["observation"] == env.transform._cat_buffers_observation).all()
         assert (
             cloned._cat_buffers_observation == env.transform._cat_buffers_observation
         ).all()
@@ -6693,6 +6692,7 @@ def _test_vecnorm_subproc_auto(
         tensordict = env.reset()
         for _ in range(10):
             tensordict = env.rand_step(tensordict)
+            tensordict = step_mdp(tensordict)
         queue_out.put(True)
         msg = queue_in.get(timeout=TIMEOUT)
         assert msg == "all_done"
@@ -6800,11 +6800,13 @@ def _run_parallelenv(parallel_env, queue_in, queue_out):
         assert msg == "start"
         for _ in range(10):
             tensordict = parallel_env.rand_step(tensordict)
+            tensordict = step_mdp(tensordict)
         queue_out.put("first round")
         msg = queue_in.get(timeout=TIMEOUT)
         assert msg == "start"
         for _ in range(10):
             tensordict = parallel_env.rand_step(tensordict)
+            tensordict = step_mdp(tensordict)
         queue_out.put("second round")
         parallel_env.close()
         queue_out.close()
@@ -6884,6 +6886,7 @@ def test_vecnorm_rollout(self, parallel, thr=0.2, N=200):
         for _ in range(N):
             td = env_t.rand_step(td)
             tds.append(td.clone())
+            td = step_mdp(td)
             if td.get("done").any():
                 td = env_t.reset()
         tds = torch.stack(tds, 0)