Repository: pytorch/rl

Issue:
[BUG] ParallelEnv does not work with PettingZooEnv.
# NOTE
Any currently-working alternatives to `ParallelEnv` are welcomed.
Please see [context](#additional-context) before suggesting `MultiSyncDataCollector`.

# Brief Description Of The Bug
Cannot create multiple `PettingZooEnv`s using `ParallelEnv`.

## To Reproduce

### Steps to reproduce the behavior:
Kudos to @matteobettini for this reproduction code.
```python 
# script.py
def base_env_fn():
        return PettingZooEnv(
            task="multiwalker_v9",
            parallel=True,
            seed=42,
            n_walkers=3,
            shared_reward=False,
            max_cycles=1000,
            render_mode=None,
            device="cpu",
        )

    collector = SyncDataCollector(
        lambda: ParallelEnv(
            num_workers=4,  # noqa: E731
            create_env_fn=base_env_fn,
            device="cpu",
        ),
        policy=None,
        frames_per_batch=100,
        max_frames_per_traj=50,
        total_frames=200,
        device="cpu",
        reset_at_each_iter=False,
    )
    try:
        for i, tensordict in enumerate(collector):
            print("Reached")
    finally:
        collector.shutdown()
```
### Error Trace
```bash
Traceback (most recent call last):
  File "/home/n00bcak/Desktop/<path_to_script>/script.py", line 32, in <module>
    for i, tensordict in enumerate(collector):
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/collectors/collectors.py", line 888, in iterator
    tensordict_out = self.rollout()
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/_utils.py", line 480, in unpack_rref_and_invoke_function
    return func(self, *args, **kwargs)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/collectors/collectors.py", line 1007, in rollout
    env_output, env_next_output = self.env.step_and_maybe_reset(env_input)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/envs/common.py", line 2750, in step_and_maybe_reset
    tensordict_ = self.maybe_reset(tensordict_)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/envs/common.py", line 2795, in maybe_reset
    tensordict = self.reset(tensordict)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/envs/common.py", line 2120, in reset
    tensordict_reset = self._reset(tensordict, **kwargs)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/envs/transforms/transforms.py", line 809, in _reset
    tensordict_reset = self.base_env._reset(tensordict, **kwargs)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/envs/batched_envs.py", line 60, in decorated_fun
    return fun(self, *args, **kwargs)
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/torchrl/envs/batched_envs.py", line 1503, in _reset
    self.shared_tensordicts[i].update_(
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/tensordict/base.py", line 3192, in update_
    self._apply_nest(
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/tensordict/_td.py", line 988, in _apply_nest
    item_trsf = item._apply_nest(
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/tensordict/_td.py", line 1007, in _apply_nest
    item_trsf = fn(
  File "/home/n00bcak/Desktop/<path_to_venv>/site-packages/tensordict/base.py", line 3176, in inplace_update
    dest.copy_(source, non_blocking=non_blocking)
RuntimeError: output with shape [3, 1] doesn't match the broadcast shape [3, 3]
```

## Expected behavior

A clear and concise description of what you expected to happen.

"Reached" is printed to `sys.stdout` AND program terminates without error.

## System info

```python
>>> import torchrl, numpy, sys
>>> print(f"TorchRL: {torchrl.__version__}\nNumPy: {numpy.__version__}\nPython3 Ver: {sys.version}\nPlatform: {sys.platform}")
TorchRL: 0.4.0 
NumPy: 1.25.0 
Python3 Ver: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] 
Platform: linux
```
```bash
> lsb_release -a
Distributor ID:	Ubuntu
Description:	Ubuntu 22.04.3 LTS
Release:	22.04
Codename:	jammy
```

## Additional context

Additional context can be found [here](https://discuss.pytorch.org/t/pettingzoo-trouble-running-multiple-marl-environments-in-parallel/203706/4).

## Reason

Some sleuthing reveals that the culprit appears to be in `torchrl/envs/utils.py` in `inner_terminated_or_truncated`:

https://github.com/pytorch/rl/blob/767a877a0a93d41a9b7935598da0ded4c984904f/torchrl/envs/utils.py#L1130-L1136

**NOTE: I use `input_data` here to refer to the input tensordict of the first `_terminated_or_truncated` call.**
Namely, this line, which erroneously checks the `aggregate.ndim` against `input_data.ndim` = `len(input_data.batch_size)`:
https://github.com/pytorch/rl/blob/767a877a0a93d41a9b7935598da0ded4c984904f/torchrl/envs/utils.py#L1132

By inserting print statements, 
```python
        if aggregate is not None:
            if key is not None:
                if aggregate.ndim > data.ndim:
                    # accounts for trailing singleton dim in done.
                    # _reset is always expanded on the right if needed so this can only be useful
                    print("BEFORE: ", data, aggregate.shape)
                    aggregate = aggregate.squeeze(-1)
                    print("AFTER: ", data, aggregate.shape, "\n\n###\n\n")
```
it may be verified that the `_reset` key is given the shape of `input_data.batch_size`.

```
###


TensorDict(
    fields={
        action: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([4, 3, 31]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=cpu,
    is_shared=False) torch.Size([4, 3, 1])
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        observation: Tensor(shape=torch.Size([4, 3, 31]), device=cpu, dtype=torch.float32, is_shared=False),
        terminated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
    batch_size=torch.Size([4, 3]),
    device=cpu,
    is_shared=False) torch.Size([4, 3]) 

###


TensorDict(
    fields={
        collector: TensorDict(
            fields={
                traj_ids: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([4]),
            device=cpu,
            is_shared=False),
        done: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        step_count: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        walker: TensorDict(
            fields={
                _reset: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.bool, is_shared=False),
                action: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                done: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([4, 3, 31]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([4, 3]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=cpu,
    is_shared=False) torch.Size([4, 1])
TensorDict(
    fields={
        collector: TensorDict(
            fields={
                traj_ids: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([4]),
            device=cpu,
            is_shared=False),
        done: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        step_count: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.int64, is_shared=False),
        terminated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        truncated: Tensor(shape=torch.Size([4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        walker: TensorDict(
            fields={
                _reset: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.bool, is_shared=False),
                action: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
                done: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([4, 3, 31]), device=cpu, dtype=torch.float32, is_shared=False),
                terminated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                truncated: Tensor(shape=torch.Size([4, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
            batch_size=torch.Size([4, 3]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=cpu,
    is_shared=False) torch.Size([4]) 

###
```

## Proposed fix

Unfortunately, as I am still new to torchRL, I cannot suggest any changes as I do not know the implications of changing this function. 

## 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:

Resolving patch:
diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
--- a/torchrl/envs/utils.py
+++ b/torchrl/envs/utils.py
@@ -1129,10 +1129,6 @@ def inner_terminated_or_truncated(data, full_done_spec, key, curr_done_key=()):
 
         if aggregate is not None:
             if key is not None:
-                if aggregate.ndim > data.ndim:
-                    # accounts for trailing singleton dim in done.
-                    # _reset is always expanded on the right if needed so this can only be useful
-                    aggregate = aggregate.squeeze(-1)
                 data.set(key, aggregate)
                 list_of_keys.append(curr_done_key + (key,))
             any_eot = any_eot | aggregate.any()


Test patch:
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
--- a/test/mocking_classes.py
+++ b/test/mocking_classes.py
@@ -1751,9 +1751,9 @@ def _reset(
         if tensordict is not None:
             _reset = tensordict.get("_reset", None)
             if _reset is not None:
-                self.count[_reset] = self.start_val
-                self.count_nested_1[_reset] = self.start_val
-                self.count_nested_2[_reset] = self.start_val
+                self.count[_reset.squeeze(-1)] = self.start_val
+                self.count_nested_1[_reset.squeeze(-1)] = self.start_val
+                self.count_nested_2[_reset.squeeze(-1)] = self.start_val
             else:
                 reset_all = True
 
diff --git a/test/test_env.py b/test/test_env.py
--- a/test/test_env.py
+++ b/test/test_env.py
@@ -2428,63 +2428,73 @@ def test_root_prevail(self, done_key):
         assert _terminated_or_truncated(data, full_done_spec=spec)
 
     def test_terminated_or_truncated_nospec(self):
-        data = TensorDict({"done": torch.zeros(2, 1, dtype=torch.bool)}, [2])
+        done_shape = (2, 1)
+        nested_done_shape = (2, 3, 1)
+        data = TensorDict(
+            {"done": torch.zeros(*done_shape, dtype=torch.bool)}, done_shape[0]
+        )
         assert not _terminated_or_truncated(data, write_full_false=True)
-        assert data["_reset"].shape == (2,)
+        assert data["_reset"].shape == done_shape
         assert not _terminated_or_truncated(data, write_full_false=False)
         assert data.get("_reset", None) is None
 
         data = TensorDict(
             {
-                ("agent", "done"): torch.zeros(2, 1, dtype=torch.bool),
-                ("nested", "done"): torch.ones(2, 1, dtype=torch.bool),
+                ("agent", "done"): torch.zeros(*nested_done_shape, dtype=torch.bool),
+                ("nested", "done"): torch.ones(*nested_done_shape, dtype=torch.bool),
             },
-            [2],
+            [done_shape[0]],
         )
         assert _terminated_or_truncated(data)
-        assert data["agent", "_reset"].shape == (2,)
-        assert data["nested", "_reset"].shape == (2,)
+        assert data["agent", "_reset"].shape == nested_done_shape
+        assert data["nested", "_reset"].shape == nested_done_shape
 
         data = TensorDict(
             {
-                "done": torch.zeros(2, 1, dtype=torch.bool),
-                ("nested", "done"): torch.zeros(2, 1, dtype=torch.bool),
+                "done": torch.zeros(*done_shape, dtype=torch.bool),
+                ("nested", "done"): torch.zeros(*nested_done_shape, dtype=torch.bool),
             },
-            [2],
+            [done_shape[0]],
         )
         assert not _terminated_or_truncated(data, write_full_false=False)
         assert data.get("_reset", None) is None
         assert data.get(("nested", "_reset"), None) is None
         assert not _terminated_or_truncated(data, write_full_false=True)
-        assert data["_reset"].shape == (2,)
-        assert data["nested", "_reset"].shape == (2,)
+        assert data["_reset"].shape == done_shape
+        assert data["nested", "_reset"].shape == nested_done_shape
 
         data = TensorDict(
             {
-                "terminated": torch.zeros(2, 1, dtype=torch.bool),
-                "truncated": torch.ones(2, 1, dtype=torch.bool),
-                ("nested", "terminated"): torch.zeros(2, 1, dtype=torch.bool),
+                "terminated": torch.zeros(*done_shape, dtype=torch.bool),
+                "truncated": torch.ones(*done_shape, dtype=torch.bool),
+                ("nested", "terminated"): torch.zeros(
+                    *nested_done_shape, dtype=torch.bool
+                ),
             },
-            [2],
+            [done_shape[0]],
         )
         assert _terminated_or_truncated(data, write_full_false=False)
-        assert data["_reset"].shape == (2,)
-        assert data["nested", "_reset"].shape == (2,)
+        assert data["_reset"].shape == done_shape
+        assert data["nested", "_reset"].shape == nested_done_shape
         assert data["_reset"].all()
         assert not data["nested", "_reset"].any()
 
     def test_terminated_or_truncated_spec(self):
+        done_shape = (2, 1)
+        nested_done_shape = (2, 3, 1)
         spec = CompositeSpec(
-            done=DiscreteTensorSpec(2, shape=(2, 1), dtype=torch.bool),
+            done=DiscreteTensorSpec(2, shape=done_shape, dtype=torch.bool),
             shape=[
                 2,
             ],
         )
-        data = TensorDict({"done": torch.zeros(2, 1, dtype=torch.bool)}, [2])
+        data = TensorDict(
+            {"done": torch.zeros(*done_shape, dtype=torch.bool)}, [done_shape[0]]
+        )
         assert not _terminated_or_truncated(
             data, write_full_false=True, full_done_spec=spec
         )
-        assert data["_reset"].shape == (2,)
+        assert data["_reset"].shape == done_shape
         assert not _terminated_or_truncated(
             data, write_full_false=False, full_done_spec=spec
         )
@@ -2493,33 +2503,31 @@ def test_terminated_or_truncated_spec(self):
         spec = CompositeSpec(
             {
                 ("agent", "done"): DiscreteTensorSpec(
-                    2, shape=(2, 1), dtype=torch.bool
+                    2, shape=nested_done_shape, dtype=torch.bool
                 ),
                 ("nested", "done"): DiscreteTensorSpec(
-                    2, shape=(2, 1), dtype=torch.bool
+                    2, shape=nested_done_shape, dtype=torch.bool
                 ),
             },
-            shape=[
-                2,
-            ],
+            shape=[nested_done_shape[0]],
         )
         data = TensorDict(
             {
-                ("agent", "done"): torch.zeros(2, 1, dtype=torch.bool),
-                ("nested", "done"): torch.ones(2, 1, dtype=torch.bool),
+                ("agent", "done"): torch.zeros(*nested_done_shape, dtype=torch.bool),
+                ("nested", "done"): torch.ones(*nested_done_shape, dtype=torch.bool),
             },
-            [2],
+            [nested_done_shape[0]],
         )
         assert _terminated_or_truncated(data, full_done_spec=spec)
-        assert data["agent", "_reset"].shape == (2,)
-        assert data["nested", "_reset"].shape == (2,)
+        assert data["agent", "_reset"].shape == nested_done_shape
+        assert data["nested", "_reset"].shape == nested_done_shape
 
         data = TensorDict(
             {
-                ("agent", "done"): torch.zeros(2, 1, dtype=torch.bool),
-                ("nested", "done"): torch.zeros(2, 1, dtype=torch.bool),
+                ("agent", "done"): torch.zeros(*nested_done_shape, dtype=torch.bool),
+                ("nested", "done"): torch.zeros(*nested_done_shape, dtype=torch.bool),
             },
-            [2],
+            [nested_done_shape[0]],
         )
         assert not _terminated_or_truncated(
             data, write_full_false=False, full_done_spec=spec
@@ -2529,32 +2537,34 @@ def test_terminated_or_truncated_spec(self):
         assert not _terminated_or_truncated(
             data, write_full_false=True, full_done_spec=spec
         )
-        assert data["agent", "_reset"].shape == (2,)
-        assert data["nested", "_reset"].shape == (2,)
+        assert data["agent", "_reset"].shape == nested_done_shape
+        assert data["nested", "_reset"].shape == nested_done_shape
 
         spec = CompositeSpec(
             {
-                "truncated": DiscreteTensorSpec(2, shape=(2, 1), dtype=torch.bool),
-                "terminated": DiscreteTensorSpec(2, shape=(2, 1), dtype=torch.bool),
+                "truncated": DiscreteTensorSpec(2, shape=done_shape, dtype=torch.bool),
+                "terminated": DiscreteTensorSpec(2, shape=done_shape, dtype=torch.bool),
                 ("nested", "terminated"): DiscreteTensorSpec(
-                    2, shape=(2, 1), dtype=torch.bool
+                    2, shape=nested_done_shape, dtype=torch.bool
                 ),
             },
             shape=[2],
         )
         data = TensorDict(
             {
-                "terminated": torch.zeros(2, 1, dtype=torch.bool),
-                "truncated": torch.ones(2, 1, dtype=torch.bool),
-                ("nested", "terminated"): torch.zeros(2, 1, dtype=torch.bool),
+                "terminated": torch.zeros(*done_shape, dtype=torch.bool),
+                "truncated": torch.ones(*done_shape, dtype=torch.bool),
+                ("nested", "terminated"): torch.zeros(
+                    *nested_done_shape, dtype=torch.bool
+                ),
             },
-            [2],
+            [done_shape[0]],
         )
         assert _terminated_or_truncated(
             data, write_full_false=False, full_done_spec=spec
         )
-        assert data["_reset"].shape == (2,)
-        assert data["nested", "_reset"].shape == (2,)
+        assert data["_reset"].shape == done_shape
+        assert data["nested", "_reset"].shape == nested_done_shape
         assert data["_reset"].all()
         assert not data["nested", "_reset"].any()
 
diff --git a/test/test_libs.py b/test/test_libs.py
--- a/test/test_libs.py
+++ b/test/test_libs.py
@@ -3518,6 +3518,32 @@ def test_vec_env(self, task, parallel, maybe_fork_ParallelEnv):
         vec_env = maybe_fork_ParallelEnv(2, create_env_fn=env_fun)
         vec_env.rollout(100, break_when_any_done=False)
 
+    def test_reset_parallel_env(self, maybe_fork_ParallelEnv):
+        def base_env_fn():
+            return PettingZooEnv(
+                task="multiwalker_v9",
+                parallel=True,
+                seed=0,
+                n_walkers=3,
+                max_cycles=1000,
+            )
+
+        collector = SyncDataCollector(
+            lambda: maybe_fork_ParallelEnv(
+                num_workers=2,
+                create_env_fn=base_env_fn,
+                device="cpu",
+            ),
+            policy=None,
+            frames_per_batch=100,
+            max_frames_per_traj=50,
+            total_frames=200,
+            reset_at_each_iter=False,
+        )
+        for _ in collector:
+            pass
+        collector.shutdown()
+
     @pytest.mark.parametrize("task", ["knights_archers_zombies_v10", "pistonball_v6"])
     @pytest.mark.parametrize("parallel", [True, False])
     def test_collector(self, task, parallel):