Repository: pytorch/rl

Issue:
[BUG] Incorrect reset handling in collectors
## Describe the bug

After auto-resetting an environment with `_reset_if_necessary`, the initial obs is ignored. The actual obs seen by the policy at the next step is always a zero TensorDict.

## To Reproduce

Here we use a dummy env where the obs is just the time stamp (starting from 1).

```python
from torchrl.collectors import SyncDataCollector
from torchrl.envs import EnvBase
from tensordict import TensorDict
from torchrl.data import TensorSpec, CompositeSpec, UnboundedContinuousTensorSpec
import torch

class DummyEnv(EnvBase):
    def __init__(
        self, 
        device = "cpu", 
        dtype = None, 
        batch_size = None, 
        run_type_checks: bool = True
    ):
        super().__init__(device, dtype, batch_size, run_type_checks)
        self.observation_spec = CompositeSpec({
            "time": UnboundedContinuousTensorSpec((*batch_size, 1)),
        }, shape=batch_size)
        self.action_spec = UnboundedContinuousTensorSpec((*batch_size, 1))
        self.reward_spec = UnboundedContinuousTensorSpec((*batch_size, 1,))
        self.time: torch.Tensor = torch.zeros(*self.batch_size, 1, device=self.device)
    
    def _step(self, tensordict: TensorDict) -> TensorDict:
        result = TensorDict({
            "time": self.time.clone(),
            "reward": self.reward_spec.rand(),
            "done": self.time > 4,
        }, self.batch_size)
        self.time += 1
        return result
    
    def _reset(self, tensordict: TensorDict, **kwargs) -> TensorDict:
        if tensordict is not None:
            reset_mask = tensordict.get("_reset")
            self.time[reset_mask] = 1
        else:
            # reset all envs
            self.time[:] = 1
        result = TensorDict({
            "time": self.time.clone()
        }, self.batch_size)
        self.time += 1
        return result

    def _set_seed(self, seed):
        torch.manual_seed(seed)

if __name__ == "__main__":
    batch_size = [4]
    env = DummyEnv(batch_size=batch_size)
    def policy(tensordict: TensorDict):
        if "collector" in tensordict.keys():
            step = tensordict[("collector", "step_count")]
            print(f'step: {step}, obs: {tensordict["time"].squeeze(-1)}')
        tensordict.set("action", env.action_spec.rand())
        return tensordict

    collector = SyncDataCollector(
        env, policy, split_trajs=False, frames_per_batch=4 * 10
    )

    for i, data in enumerate(collector):
        break

```

Running the above code gives

```
step: tensor([0, 0, 0, 0], dtype=torch.int32), obs: tensor([1., 1., 1., 1.]) # the first episode is correct
step: tensor([1, 1, 1, 1], dtype=torch.int32), obs: tensor([2., 2., 2., 2.])
step: tensor([2, 2, 2, 2], dtype=torch.int32), obs: tensor([3., 3., 3., 3.])
step: tensor([3, 3, 3, 3], dtype=torch.int32), obs: tensor([4., 4., 4., 4.])
step: tensor([0, 0, 0, 0], dtype=torch.int32), obs: tensor([0., 0., 0., 0.]) # the policy should have never seen zeros
step: tensor([1, 1, 1, 1], dtype=torch.int32), obs: tensor([2., 2., 2., 2.])
step: tensor([2, 2, 2, 2], dtype=torch.int32), obs: tensor([3., 3., 3., 3.])
step: tensor([3, 3, 3, 3], dtype=torch.int32), obs: tensor([4., 4., 4., 4.])
step: tensor([0, 0, 0, 0], dtype=torch.int32), obs: tensor([0., 0., 0., 0.]) # the policy should have never seen zeros
step: tensor([1, 1, 1, 1], dtype=torch.int32), obs: tensor([2., 2., 2., 2.])
```

## Expected behavior

```
step: tensor([0, 0, 0, 0], dtype=torch.int32), obs: tensor([1., 1., 1., 1.]) # the first episode is correct
step: tensor([1, 1, 1, 1], dtype=torch.int32), obs: tensor([2., 2., 2., 2.])
step: tensor([2, 2, 2, 2], dtype=torch.int32), obs: tensor([3., 3., 3., 3.])
step: tensor([3, 3, 3, 3], dtype=torch.int32), obs: tensor([4., 4., 4., 4.])
step: tensor([0, 0, 0, 0], dtype=torch.int32), obs: tensor([1., 1., 1., 1.]) # consistent with the first episode
step: tensor([1, 1, 1, 1], dtype=torch.int32), obs: tensor([2., 2., 2., 2.])
step: tensor([2, 2, 2, 2], dtype=torch.int32), obs: tensor([3., 3., 3., 3.])
step: tensor([3, 3, 3, 3], dtype=torch.int32), obs: tensor([4., 4., 4., 4.])
step: tensor([0, 0, 0, 0], dtype=torch.int32), obs: tensor([1., 1., 1., 1.]) # consistent with the first episode
step: tensor([1, 1, 1, 1], dtype=torch.int32), obs: tensor([2., 2., 2., 2.])
```

## Reason and Possible fixes

This occurs because in `SyncDataCollector`:

```python
def rollout(self):
    ...
    self._reset_if_necessary()
    self._tensordict.update(step_mdp(self._tensordict), inplace=True) 
    ...

def _reset_if_necessary(self):
    ...
    if done_or_terminated.any():
        traj_ids = self._tensordict.get(("collector", "traj_ids")).clone()
        steps = steps.clone()
        if len(self.env.batch_size):
            self._tensordict.masked_fill_(done_or_terminated, 0)
            _reset = done_or_terminated
            self._tensordict.set("_reset", _reset)
        else:
            _reset = None
            self._tensordict.zero_()
        self.env.reset(self._tensordict)
    ...
```

The initial obs of the new episode gets discarded by `step_mdp` because it is not in `self._tensordict["next"]`. What the policy will see is the zeros set by `self._tensordict.masked_fill_(done_or_terminated, 0)`.

The most straightforward fix is to change the above to:

```python
def _reset_if_necessary(self):
    ...
    if done_or_terminated.any():
        traj_ids = self._tensordict.get(("collector", "traj_ids")).clone()
        steps = steps.clone()
        if len(self.env.batch_size):
            self._tensordict.masked_fill_(done_or_terminated, 0)
            _reset = done_or_terminated
            self._tensordict.set(("next", "_reset"), _reset)
        else:
            _reset = None
            self._tensordict.zero_()
        self.env.reset(self._tensordict["next"])
    ...
```
So that the initial obs get carried to the next step by `step_mdp`.

However, this would break some tests, e.g., `test_collector.py::test_traj_len_consistency`, because now we have `("next", "done")` in keys after the first reset which causes key inconsistency when doing `torch.cat`. 

I recall that the earlier versions of torchrl require `step_mdp(env.reset())`.  Here I wonder does the coexistence of "done" and ("next", "done") make sense. I personally think having both is more rigorous: "done" indicates whether this step is an initial step, and ("next", "done") indicates whether the episode is terminated after this step (IIUC currently we have only "done" for the latter). In this way inside an RNN policy module, we can decide whether we need to reset some of its hidden states. 

## 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:
thanks for reporting this.
I think we should first fix the bug and then consider having ("next", "done") along with "done".
It's a challenging one: there are some conventions, like `r_t` goes with `s_t`, `a_t`, only `s_{t+1}` is considered as the "next" value. Is the related `final_t` a property of `t` or `t+1` is not super clear IMO.
IIUC what you're suggesting, we should move the current done to the next state, is that right?
I'm just afraid that it would not be obvious for all users. But I'm open to discussing it...

I'll make a PR to fix this bug but feel free to keep the conversation going regarding `"done"`.

Resolving patch:
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
--- a/torchrl/collectors/collectors.py
+++ b/torchrl/collectors/collectors.py
@@ -571,7 +571,7 @@ def iterator(self) -> Iterator[TensorDictBase]:
             if self._frames >= self.total_frames:
                 break
 
-    def _reset_if_necessary(self) -> None:
+    def _step_and_maybe_reset(self) -> None:
         done = self._tensordict.get("done")
         if not self.reset_when_done:
             done = torch.zeros_like(done)
@@ -592,6 +592,11 @@ def _reset_if_necessary(self) -> None:
             done_or_terminated = done_or_terminated | _reset
 
         if done_or_terminated.any():
+            if not done_or_terminated.all():
+                self._tensordict[~done_or_terminated] = step_mdp(
+                    self._tensordict[~done_or_terminated]
+                )
+
             traj_ids = self._tensordict.get(("collector", "traj_ids")).clone()
             steps = steps.clone()
             if len(self.env.batch_size):
@@ -617,6 +622,8 @@ def _reset_if_necessary(self) -> None:
                 ("collector", "traj_ids"), traj_ids
             )  # no ops if they already match
             self._tensordict.set_(("collector", "step_count"), steps)
+        else:
+            self._tensordict.update(step_mdp(self._tensordict), inplace=True)
 
     @torch.no_grad()
     def rollout(self) -> TensorDictBase:
@@ -651,8 +658,7 @@ def rollout(self) -> TensorDictBase:
                     if is_shared:
                         self._tensordict_out.share_memory_()
 
-                self._reset_if_necessary()
-                self._tensordict.update(step_mdp(self._tensordict), inplace=True)
+                self._step_and_maybe_reset()
 
         return self._tensordict_out
 


Test patch:
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
--- a/test/mocking_classes.py
+++ b/test/mocking_classes.py
@@ -97,6 +97,8 @@ def custom_td(self):
 
 
 class MockSerialEnv(EnvBase):
+    """A simple counting env that is reset after a predifined max number of steps."""
+
     @classmethod
     def __new__(
         cls,
@@ -844,9 +846,16 @@ def forward(self, observation, action):
 
 
 class CountingEnv(EnvBase):
-    def __init__(self, max_steps: int = 5, **kwargs):
+    """An env that is done after a given number of steps.
+
+    The action is the count increment.
+
+    """
+
+    def __init__(self, max_steps: int = 5, start_val: int = 0, **kwargs):
         super().__init__(**kwargs)
         self.max_steps = max_steps
+        self.start_val = start_val
 
         self.observation_spec = CompositeSpec(
             observation=UnboundedContinuousTensorSpec(
@@ -878,9 +887,9 @@ def _set_seed(self, seed: Optional[int]):
     def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
         if tensordict is not None and "_reset" in tensordict.keys():
             _reset = tensordict.get("_reset")
-            self.count[_reset] = 0
+            self.count[_reset] = self.start_val
         else:
-            self.count[:] = 0
+            self.count[:] = self.start_val
         return TensorDict(
             source={
                 "observation": self.count.clone(),
@@ -905,3 +914,87 @@ def _step(
             batch_size=self.batch_size,
             device=self.device,
         )
+
+
+class CountingBatchedEnv(EnvBase):
+    """An env that is done after a given number of steps.
+
+    The action is the count increment.
+
+    Unlike ``CountingEnv``, different envs of the batch can have different max_steps
+    """
+
+    def __init__(
+        self,
+        max_steps: torch.Tensor = None,
+        start_val: torch.Tensor = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        if max_steps is None:
+            max_steps = torch.tensor(5)
+        if start_val is None:
+            start_val = torch.zeros(())
+        if not max_steps.shape == self.batch_size:
+            raise RuntimeError("batch_size and max_steps shape must match.")
+
+        self.max_steps = max_steps
+        self.start_val = start_val
+
+        self.observation_spec = CompositeSpec(
+            observation=UnboundedContinuousTensorSpec(
+                (
+                    *self.batch_size,
+                    1,
+                )
+            ),
+            shape=self.batch_size,
+        )
+        self.reward_spec = UnboundedContinuousTensorSpec(
+            (
+                *self.batch_size,
+                1,
+            )
+        )
+        self.input_spec = CompositeSpec(
+            action=BinaryDiscreteTensorSpec(n=1, shape=[*self.batch_size, 1]),
+            shape=self.batch_size,
+        )
+
+        self.count = torch.zeros(
+            (*self.batch_size, 1), device=self.device, dtype=torch.int
+        )
+
+    def _set_seed(self, seed: Optional[int]):
+        torch.manual_seed(seed)
+
+    def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
+        if tensordict is not None and "_reset" in tensordict.keys():
+            _reset = tensordict.get("_reset")
+            self.count[_reset] = self.start_val[_reset].unsqueeze(-1)
+        else:
+            self.count[:] = self.start_val.unsqueeze(-1)
+        return TensorDict(
+            source={
+                "observation": self.count.clone(),
+                "done": self.count > self.max_steps.unsqueeze(-1),
+            },
+            batch_size=self.batch_size,
+            device=self.device,
+        )
+
+    def _step(
+        self,
+        tensordict: TensorDictBase,
+    ) -> TensorDictBase:
+        action = tensordict.get("action")
+        self.count += action.to(torch.int).unsqueeze(-1)
+        return TensorDict(
+            source={
+                "observation": self.count,
+                "done": self.count > self.max_steps.unsqueeze(-1),
+                "reward": torch.zeros_like(self.count, dtype=torch.float),
+            },
+            batch_size=self.batch_size,
+            device=self.device,
+        )
diff --git a/test/test_collector.py b/test/test_collector.py
--- a/test/test_collector.py
+++ b/test/test_collector.py
@@ -12,6 +12,8 @@
 from _utils_internal import generate_seeds, PENDULUM_VERSIONED, PONG_VERSIONED
 from mocking_classes import (
     ContinuousActionVecMockEnv,
+    CountingBatchedEnv,
+    CountingEnv,
     DiscreteActionConvMockEnv,
     DiscreteActionConvPolicy,
     DiscreteActionVecMockEnv,
@@ -1181,6 +1183,52 @@ def test_auto_wrap_error(self, collector_class, env_maker):
             )
 
 
+@pytest.mark.parametrize("env_class", [CountingEnv, CountingBatchedEnv])
+def test_initial_obs_consistency(env_class, seed=1):
+    torch.manual_seed(seed)
+    start_val = 4
+    if env_class == CountingEnv:
+        num_envs = 1
+        env = CountingEnv(device="cpu", max_steps=8, start_val=start_val)
+        max_steps = 8
+    elif env_class == CountingBatchedEnv:
+        num_envs = 2
+        env = CountingBatchedEnv(
+            device="cpu",
+            batch_size=[num_envs],
+            max_steps=torch.arange(num_envs) + 17,
+            start_val=torch.ones([num_envs]) * start_val,
+        )
+        max_steps = env.max_steps.max().item()
+    env.set_seed(seed)
+    policy = lambda tensordict: tensordict.set(
+        "action", torch.ones(tensordict.shape, dtype=torch.int)
+    )
+    collector = SyncDataCollector(
+        create_env_fn=env,
+        policy=policy,
+        frames_per_batch=((max_steps - 3) * 2 + 2) * num_envs,  # at least two episodes
+        split_trajs=False,
+    )
+    for _d in collector:
+        break
+    obs = _d["observation"].squeeze()
+    if env_class == CountingEnv:
+        arange_0 = start_val + torch.arange(max_steps - 3)
+        arange = start_val + torch.arange(2)
+        expected = torch.cat([arange_0, arange_0, arange]).float()
+    else:
+        # the first env has a shorter horizon than the second
+        arange_0 = start_val + torch.arange(max_steps - 3 - 1)
+        arange = start_val + torch.arange(start_val)
+        expected_0 = torch.cat([arange_0, arange_0, arange]).float()
+        arange_0 = start_val + torch.arange(max_steps - 3)
+        arange = start_val + torch.arange(2)
+        expected_1 = torch.cat([arange_0, arange_0, arange]).float()
+        expected = torch.stack([expected_0, expected_1])
+    assert torch.allclose(obs, expected)
+
+
 def weight_reset(m):
     if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
         m.reset_parameters()