Repository: pytorch/rl

Issue:
[BUG] data collectors may omit some keys in the output tensordict
## Describe the bug

Remote data collector will populate the output tensordict with the data that come with the first collection in the environment. There are several issues with this, the main being that if we use a value of `init_random_frames` > 0, the initial tensordict will just be made out of the environment keys + an action key + some utils from the collector (trajectory id and such).

## To Reproduce

We can create a policy that uses an RNN to reproduce this.
```python
from torchrl.collectors import MultiaSyncDataCollector
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import TensorDictModule, LSTMNet
import torch

out_features=1
batch = 7
time_steps = 6
hidden_size = 12
net = LSTMNet(
    out_features,
    {"input_size": hidden_size, "hidden_size": hidden_size},
    {"out_features": hidden_size},
)

policy = TensorDictModule(net, ["observation", "hidden1", "hidden2"], ["action", "hidden1", "hidden2", "next_hidden1", "next_hidden2"])

env_maker = lambda: GymEnv("Pendulum-v1")

policy(env_maker().reset())
```

### Case 1: no `init_random_frames`

Here everything works as expected
```python
c = MultiaSyncDataCollector([env_maker], policy, total_frames=200, frames_per_batch=20, init_random_frames=0)

for i, b in enumerate(c):
    print(b)
    if i == 5:
        # by now we should have a tensordict that is the output of the policy
        break
```
Every tensordict look like this

```
TensorDict(
    fields={
        action: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        done: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        hidden1: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        hidden2: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        mask: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        next_hidden1: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        next_hidden2: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        next_observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        reward: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        step_count: Tensor(torch.Size([1, 20, 1]), dtype=torch.int32),
        traj_ids: Tensor(torch.Size([1, 20, 1, 1]), dtype=torch.int64)},
    batch_size=torch.Size([1, 20]),
    device=cpu,
    is_shared=False)
```

### Case 2: `init_random_frames` > 0

When running this
```python
c = MultiaSyncDataCollector([env_maker], policy, total_frames=200, frames_per_batch=20, init_random_frames=50)

for i, b in enumerate(c):
    print(b)
    if i == 5:
        # by now we should have a tensordict that is the output of the policy
        break
```
every tensordict look like this

```
TensorDict(
    fields={
        action: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        done: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        mask: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        next_observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        reward: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        step_count: Tensor(torch.Size([1, 20, 1]), dtype=torch.int32),
        traj_ids: Tensor(torch.Size([1, 20, 1, 1]), dtype=torch.int64)},
    batch_size=torch.Size([1, 20]),
    device=cpu,
    is_shared=False)
```
i.e. the `hidden` and `next_hidden` keys are missing **even when we have passed the number of frames indicated by `init_random_frames`**.

### Case 3: `init_random_frames` > 0 with `SyncDataCollector`

`SyncDataCollector` is the simplest data collector: it does not spawn other processes. 

```python
c = SyncDataCollector(env_maker, policy, total_frames=200, frames_per_batch=20, init_random_frames=50)

for i, b in enumerate(c):
    print(b)
    if i == 5:
        # by now we should have a tensordict that is the output of the policy
        break
```
The first three tensordicts look like this
```
TensorDict(
    fields={
        action: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        done: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        mask: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        next_observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        reward: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        step_count: Tensor(torch.Size([1, 20, 1]), dtype=torch.int32),
        traj_ids: Tensor(torch.Size([1, 20, 1, 1]), dtype=torch.int64)},
    batch_size=torch.Size([1, 20]),
    device=cpu,
    is_shared=False)
```
after that, they start being populated with the hidden keys:

```
TensorDict(
    fields={
        action: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        done: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        hidden1: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        hidden2: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        mask: Tensor(torch.Size([1, 20, 1]), dtype=torch.bool),
        next_hidden1: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        next_hidden2: Tensor(torch.Size([1, 20, 1, 12]), dtype=torch.float32),
        next_observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        observation: Tensor(torch.Size([1, 20, 3]), dtype=torch.float32),
        reward: Tensor(torch.Size([1, 20, 1]), dtype=torch.float32),
        step_count: Tensor(torch.Size([1, 20, 1]), dtype=torch.int32),
        traj_ids: Tensor(torch.Size([1, 20, 1, 1]), dtype=torch.int64)},
    batch_size=torch.Size([1, 20]),
    device=cpu,
    is_shared=False)
```


## Proposed solution

The expected behaviour should be that every tensordict has all possible keys **from the first batch** (values that are not valid should be populated with 0s) for all data collectors (note that under the hood every collector will call `SyncDataCollector` hence only this class needs to be corrected).

The reason for this is that we will store the tensordict in a replay buffer and we need to know what key-value pairs will need to be stored in it from the first collection.

What we want is that after instantiating the env, the `self._tensordict_out` object should be populated with all the relevant keys.

The collector should make the following operations:
1. If the policy has an non-empty `policy.spec` attribute **and** all the keys are not None (`policy.spec` is a `CompositeSpec` object that can have values set to None) **and** the keys in the spec match the `out_keys` of the policy,  then we assume that the user has given all the relevant information about it. We just need to do 
```python
self._tensordict_out = env.reset()
self._tensordict_out.update(policy.action_spec.rand())
```
If some values in the `spec` are None or if the policy `out_keys` are not all present in `policy.spec`, we would revert to (2) here below.

2. If the policy does not have a spec, then we should run a mini-rollout in the env with the policy to make sure that we have all the relevant keys:
```python
self._tensordict_out = env.rollout(3, policy)[..., 0]
```

Finally, if at some point there are more keys in the tensordict collected in `SyncDataCollector.rollout` than in the `SyncDataCollector._tensordict_out` object we should raise a warning to the user informing that some data is likely to be missed during collection for the reasons explained above.

```


Discussion:

Resolving patch:
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
--- a/torchrl/collectors/collectors.py
+++ b/torchrl/collectors/collectors.py
@@ -322,10 +322,62 @@ def __init__(
         self._tensordict.set(
             "step_count", torch.zeros(*self.env.batch_size, 1, dtype=torch.int)
         )
-        self._tensordict_out = TensorDict(
-            {},
-            batch_size=[*self.env.batch_size, self.frames_per_batch],
-            device=self.passing_device,
+
+        if (
+            hasattr(policy, "spec")
+            and policy.spec is not None
+            and all(v is not None for v in policy.spec.values())
+            and set(policy.spec.keys()) == set(policy.out_keys)
+        ):
+            # if policy spec is non-empty, all the values are not None and the keys
+            # match the out_keys we assume the user has given all relevant information
+            self._tensordict_out = TensorDict(
+                {
+                    **env.observation_spec.zero(env.batch_size),
+                    "reward": env.reward_spec.zero(env.batch_size),
+                    "done": torch.zeros(
+                        env.batch_size, dtype=torch.bool, device=env.device
+                    ),
+                    **policy.spec.zero(env.batch_size),
+                },
+                env.batch_size,
+                device=env.device,
+            )
+            self._tensordict_out = (
+                self._tensordict_out.unsqueeze(-1)
+                .expand(*env.batch_size, self.frames_per_batch)
+                .to_tensordict()
+            )
+            self._tensordict_out = self._tensordict_out.update(
+                step_mdp(self._tensordict_out)
+            )  # add "observation" when there is "next_observation"
+        else:
+            # otherwise, we perform a small number of steps with the policy to
+            # determine the relevant keys with which to pre-populate _tensordict_out.
+            # See #505 for additional context.
+            self._tensordict_out = env.rollout(3, policy)
+            if env.batch_size:
+                self._tensordict_out = self._tensordict_out[..., :1]
+            else:
+                self._tensordict_out = self._tensordict_out[:1]
+            self._tensordict_out = (
+                self._tensordict_out.expand(*env.batch_size, self.frames_per_batch)
+                .to_tensordict()
+                .zero_()
+                .detach()
+            )
+            env.reset()
+
+        # in addition to outputs of the policy, we add traj_ids and step_count to
+        # _tensordict_out which will be collected during rollout
+        if len(self.env.batch_size):
+            traj_ids = torch.zeros(*self._tensordict_out.batch_size, 1)
+        else:
+            traj_ids = torch.zeros(*self._tensordict_out.batch_size, 1, 1)
+
+        self._tensordict_out.set("traj_ids", traj_ids)
+        self._tensordict_out.set(
+            "step_count", torch.zeros(*self._tensordict_out.batch_size, 1)
         )
 
         self.return_in_place = return_in_place


Test patch:
diff --git a/test/test_collector.py b/test/test_collector.py
--- a/test/test_collector.py
+++ b/test/test_collector.py
@@ -25,11 +25,17 @@
     MultiSyncDataCollector,
     MultiaSyncDataCollector,
 )
+from torchrl.data import (
+    CompositeSpec,
+    NdUnboundedContinuousTensorSpec,
+    UnboundedContinuousTensorSpec,
+)
 from torchrl.data.tensordict.tensordict import assert_allclose_td
 from torchrl.envs import EnvCreator
 from torchrl.envs import ParallelEnv
 from torchrl.envs.libs.gym import _has_gym
 from torchrl.envs.transforms import TransformedEnv, VecNorm
+from torchrl.modules import LSTMNet, TensorDictModule
 from torchrl.modules import OrnsteinUhlenbeckProcessWrapper, Actor
 
 # torch.set_default_dtype(torch.double)
@@ -673,14 +679,22 @@ def test_collector_vecnorm_envcreator(static_seed):
 
 
 @pytest.mark.parametrize("use_async", [False, True])
-@pytest.mark.skipif(torch.cuda.device_count() <= 1, reason="no cuda device found")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="no cuda device found")
 def test_update_weights(use_async):
-    policy = torch.nn.Linear(3, 4).cuda(1)
+    def create_env():
+        return ContinuousActionVecMockEnv()
+
+    n_actions = ContinuousActionVecMockEnv().action_spec.shape[-1]
+    policy = TensorDictModule(
+        torch.nn.LazyLinear(n_actions), in_keys=["observation"], out_keys=["action"]
+    )
+    policy(create_env().reset())
+
     collector_class = (
         MultiSyncDataCollector if not use_async else MultiaSyncDataCollector
     )
     collector = collector_class(
-        [lambda: DiscreteActionVecMockEnv()] * 3,
+        [create_env] * 3,
         policy=policy,
         devices=[torch.device("cuda:0")] * 3,
         passing_devices=[torch.device("cuda:0")] * 3,
@@ -769,6 +783,89 @@ def make_env():
     dummy_env.close()
 
 
+@pytest.mark.skipif(not _has_gym, reason="test designed with GymEnv")
+@pytest.mark.parametrize(
+    "collector_class",
+    [
+        SyncDataCollector,
+        MultiaSyncDataCollector,
+        MultiSyncDataCollector,
+    ],
+)
+@pytest.mark.parametrize("init_random_frames", [0, 50])
+@pytest.mark.parametrize("explicit_spec", [True, False])
+def test_collector_output_keys(collector_class, init_random_frames, explicit_spec):
+    from torchrl.envs.libs.gym import GymEnv
+
+    out_features = 1
+    hidden_size = 12
+    total_frames = 200
+    frames_per_batch = 20
+    num_envs = 3
+
+    net = LSTMNet(
+        out_features,
+        {"input_size": hidden_size, "hidden_size": hidden_size},
+        {"out_features": hidden_size},
+    )
+
+    policy_kwargs = {
+        "module": net,
+        "in_keys": ["observation", "hidden1", "hidden2"],
+        "out_keys": ["action", "hidden1", "hidden2", "next_hidden1", "next_hidden2"],
+    }
+    if explicit_spec:
+        hidden_spec = NdUnboundedContinuousTensorSpec((1, hidden_size))
+        policy_kwargs["spec"] = CompositeSpec(
+            action=UnboundedContinuousTensorSpec(),
+            hidden1=hidden_spec,
+            hidden2=hidden_spec,
+            next_hidden1=hidden_spec,
+            next_hidden2=hidden_spec,
+        )
+
+    policy = TensorDictModule(**policy_kwargs)
+
+    env_maker = lambda: GymEnv("Pendulum-v1")
+
+    policy(env_maker().reset())
+
+    collector_kwargs = {
+        "create_env_fn": env_maker,
+        "policy": policy,
+        "total_frames": total_frames,
+        "frames_per_batch": frames_per_batch,
+        "init_random_frames": init_random_frames,
+    }
+
+    if collector_class is not SyncDataCollector:
+        collector_kwargs["create_env_fn"] = [
+            collector_kwargs["create_env_fn"] for _ in range(num_envs)
+        ]
+
+    collector = collector_class(**collector_kwargs)
+
+    keys = [
+        "action",
+        "done",
+        "hidden1",
+        "hidden2",
+        "mask",
+        "next_hidden1",
+        "next_hidden2",
+        "next_observation",
+        "observation",
+        "reward",
+        "step_count",
+        "traj_ids",
+    ]
+    b = next(iter(collector))
+
+    assert set(b.keys()) == set(keys)
+    collector.shutdown()
+    del collector
+
+
 def weight_reset(m):
     if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
         m.reset_parameters()