Repository: pytorch/rl

Issue:
[BUG] SyncDataCollector resets all workers instead of one
## Describe the bug

When using the SyncDataCollector with a ParallelEnv, **all** workers are reset if any of the workers is done.

## To Reproduce

```python
import torch
from torchrl.collectors import SyncDataCollector
from torchrl.envs import ParallelEnv
from torchrl.envs.libs.gym import GymEnv

if __name__ == '__main__':
    num_workers = 2
    num_steps = 4000

    create_env = lambda: GymEnv(env_name='PongNoFrameskip-v4', frame_skip=4)
    create_parallel_env = lambda: ParallelEnv(num_workers, create_env)

    collector = SyncDataCollector(create_parallel_env, frames_per_batch=num_workers, split_trajs=False)
    iterator = collector.iterator()
    episode_rewards = []
    cumulative_reward = torch.zeros(num_workers)
    for _ in range(num_steps):
        batch = next(iterator).squeeze(1)
        cumulative_reward += batch['reward'].squeeze(-1)
        done = batch['done'].squeeze(-1)
        episode_rewards.extend(cumulative_reward[done].tolist())
        cumulative_reward[done] = 0.
    collector.shutdown()
    print(episode_rewards)
```

This leads to the following episode rewards, although the lowest value for Pong should be -21:
```bash
[-21.0, -37.0, -41.0, -21.0]
```
The reason is that some episodes are reset without notice.

## Expected behavior

Only the workers which are actually done should be reset.

## Reason and Possible fixes

[In this line](https://github.com/pytorch/rl/blob/4862459866153a499a3c7e5d6df4238ffdc29fe8/torchrl/collectors/collectors.py#L588) the ParallelEnv is reset but `"reset_workers"` is ignored.  
A quick fix would be to replace the line with:
```python
self._tensordict.update(self.env.reset(self._tensordict), inplace=True)
```
This leads to the following output for the above code and seems to be correct:
```bash
[-21.0, -21.0, -21.0, -21.0, -21.0, -20.0, -20.0, -21.0]
```
But maybe there is a better solution.

## 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**)

[BUG] SyncDataCollector resets all workers instead of one
## Describe the bug

When using the SyncDataCollector with a ParallelEnv, **all** workers are reset if any of the workers is done.

## To Reproduce

```python
import torch
from torchrl.collectors import SyncDataCollector
from torchrl.envs import ParallelEnv
from torchrl.envs.libs.gym import GymEnv

if __name__ == '__main__':
    num_workers = 2
    num_steps = 4000

    create_env = lambda: GymEnv(env_name='PongNoFrameskip-v4', frame_skip=4)
    create_parallel_env = lambda: ParallelEnv(num_workers, create_env)

    collector = SyncDataCollector(create_parallel_env, frames_per_batch=num_workers, split_trajs=False)
    iterator = collector.iterator()
    episode_rewards = []
    cumulative_reward = torch.zeros(num_workers)
    for _ in range(num_steps):
        batch = next(iterator).squeeze(1)
        cumulative_reward += batch['reward'].squeeze(-1)
        done = batch['done'].squeeze(-1)
        episode_rewards.extend(cumulative_reward[done].tolist())
        cumulative_reward[done] = 0.
    collector.shutdown()
    print(episode_rewards)
```

This leads to the following episode rewards, although the lowest value for Pong should be -21:
```bash
[-21.0, -37.0, -41.0, -21.0]
```
The reason is that some episodes are reset without notice.

## Expected behavior

Only the workers which are actually done should be reset.

## Reason and Possible fixes

[In this line](https://github.com/pytorch/rl/blob/4862459866153a499a3c7e5d6df4238ffdc29fe8/torchrl/collectors/collectors.py#L588) the ParallelEnv is reset but `"reset_workers"` is ignored.  
A quick fix would be to replace the line with:
```python
self._tensordict.update(self.env.reset(self._tensordict), inplace=True)
```
This leads to the following output for the above code and seems to be correct:
```bash
[-21.0, -21.0, -21.0, -21.0, -21.0, -20.0, -20.0, -21.0]
```
But maybe there is a better solution.

## 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:
Good catch I wonder why we haven't spotted this earlier!
The solution seems legit.
You can make a PR out of it, otherwise I'll take care of it.

Thanks for the thorough description 
Good catch I wonder why we haven't spotted this earlier!
The solution seems legit.
You can make a PR out of it, otherwise I'll take care of it.

Thanks for the thorough description 

Resolving patch:
diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
--- a/torchrl/collectors/collectors.py
+++ b/torchrl/collectors/collectors.py
@@ -585,7 +585,7 @@ def _reset_if_necessary(self) -> None:
             else:
                 self._tensordict.zero_()
 
-            self._tensordict.update(self.env.reset(), inplace=True)
+            self.env.reset(self._tensordict)
             if self._tensordict.get("done").any():
                 raise RuntimeError(
                     f"Got {sum(self._tensordict.get('done'))} done envs after reset."


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,6 +25,7 @@
     MultiSyncDataCollector,
     RandomPolicy,
 )
+from torchrl.collectors.utils import split_trajectories
 from torchrl.data import (
     CompositeSpec,
     NdUnboundedContinuousTensorSpec,
@@ -32,8 +33,8 @@
     UnboundedContinuousTensorSpec,
 )
 from torchrl.data.tensordict.tensordict import assert_allclose_td
-from torchrl.envs import EnvCreator, ParallelEnv
-from torchrl.envs.libs.gym import _has_gym
+from torchrl.envs import EnvCreator, ParallelEnv, SerialEnv
+from torchrl.envs.libs.gym import _has_gym, GymEnv
 from torchrl.envs.transforms import TransformedEnv, VecNorm
 from torchrl.modules import (
     Actor,
@@ -295,6 +296,32 @@ def env_fn(seed):
     ccollector.shutdown()
 
 
+@pytest.mark.skipif(not _has_gym, reason="gym library is not installed")
+def test_collector_env_reset():
+    torch.manual_seed(0)
+    env = SerialEnv(2, lambda: GymEnv("ALE/Pong-v5", frame_skip=4))
+    # env = SerialEnv(3, lambda: GymEnv("CartPole-v1", frame_skip=4))
+    env.set_seed(0)
+    collector = SyncDataCollector(
+        env, total_frames=10000, frames_per_batch=10000, split_trajs=False
+    )
+    for data in collector:
+        continue
+    steps = data["step_count"][..., 1:, :]
+    done = data["done"][..., :-1, :]
+    # we don't want just one done
+    assert done.sum() > 3
+    # check that after a done, the next step count is always 1
+    assert (steps[done] == 1).all()
+    # check that if the env is not done, the next step count is > 1
+    assert (steps[~done] > 1).all()
+    # check that if step is 1, then the env was done before
+    assert (steps == 1)[done].all()
+    # check that split traj has a minimum total reward of -21 (for pong only)
+    data = split_trajectories(data)
+    assert data["reward"].sum(-2).min() == -21
+
+
 @pytest.mark.parametrize("num_env", [1, 3])
 @pytest.mark.parametrize("env_name", ["vec"])
 def test_collector_done_persist(num_env, env_name, seed=5):