Repository: pytorch/rl

Issue:
[BUG] Truncated key is on different devices with BatchedEnv on different device than single envs.
## To Reproduce

```python
from tensordict.nn import TensorDictModule
from torch import nn
from torchrl.envs import ExplorationType, SerialEnv
from torchrl.envs import StepCounter, TransformedEnv, ParallelEnv, EnvCreator
from torchrl.envs.libs.gym import GymEnv
from torchrl.modules import ProbabilisticActor, OneHotCategorical

max_step = 200
env_id = "CartPole-v1"
device = "cuda:0"


def build_single_env():
    env = GymEnv(env_id)
    env = TransformedEnv(env)
    env.append_transform(StepCounter(max_steps=max_step))
    return env


def build_actor(env):
    return ProbabilisticActor(
        module=TensorDictModule(
            nn.LazyLinear(env.action_spec.space.n),
            in_keys=["observation"],
            out_keys=["logits"],
        ),
        spec=env.action_spec,
        distribution_class=OneHotCategorical,
        in_keys=["logits"],
        default_interaction_type=ExplorationType.RANDOM,
    )


if __name__ == "__main__":
    env = ParallelEnv(4, EnvCreator(lambda: build_single_env()), device=device)
    env = TransformedEnv(env)
    # Comment the line above to remove the problem.

    policy_module = build_actor(env)
    policy_module.to(device)
    policy_module(env.reset())
    for i in range(100):
        batches = env.rollout(max_step + 3, policy=policy_module, break_when_any_done=False)

```

```bash
/usr/bin/python3.10 /claire-rcp-scratch/home/moalla/torchrl/tests/issue_truncated_device.py 
/usr/local/lib/python3.10/dist-packages/gymnasium/core.py:311: UserWarning: WARN: env.num_envs to get variables from other wrappers is deprecated and will be removed in v1.0, to get this variable you can do `env.unwrapped.num_envs` for environment variables or `env.get_wrapper_attr('num_envs')` that will search the reminding wrappers.
  logger.warn(
/usr/local/lib/python3.10/dist-packages/gymnasium/core.py:311: UserWarning: WARN: env.reward_space to get variables from other wrappers is deprecated and will be removed in v1.0, to get this variable you can do `env.unwrapped.reward_space` for environment variables or `env.get_wrapper_attr('reward_space')` that will search the reminding wrappers.
  logger.warn(
/usr/local/lib/python3.10/dist-packages/torch/nn/modules/lazy.py:180: UserWarning: Lazy modules are a new feature under heavy development so changes to the API or functionality can happen at any moment.
  warnings.warn('Lazy modules are a new feature under heavy development '
Process _ProcessNoWarn-1:
Traceback (most recent call last):
  File "/usr/lib/python3.10/multiprocessing/process.py", line 314, in _bootstrap
    self.run()
  File "/claire-rcp-scratch/home/moalla/torchrl/repos/rl/torchrl/_utils.py", line 654, in run
    return mp.Process.run(self, *args, **kwargs)
  File "/usr/lib/python3.10/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/claire-rcp-scratch/home/moalla/torchrl/repos/rl/torchrl/envs/batched_envs.py", line 1499, in _run_worker_pipe_shared_mem
    cur_td = env.reset(tensordict=data)
  File "/claire-rcp-scratch/home/moalla/torchrl/repos/rl/torchrl/envs/common.py", line 2056, in reset
    tensordict_reset = self._reset(tensordict, **kwargs)
  File "/claire-rcp-scratch/home/moalla/torchrl/repos/rl/torchrl/envs/transforms/transforms.py", line 790, in _reset
    tensordict_reset = self.transform._reset(tensordict, tensordict_reset)
  File "/claire-rcp-scratch/home/moalla/torchrl/repos/rl/torchrl/envs/transforms/transforms.py", line 1050, in _reset
    tensordict_reset = t._reset(tensordict, tensordict_reset)
  File "/claire-rcp-scratch/home/moalla/torchrl/repos/rl/torchrl/envs/transforms/transforms.py", line 5156, in _reset
    truncated = truncated | tensordict_reset.get(truncated_key, False)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```


## System info

```python
import torchrl, tensordict, torch, numpy, sys
print(torch.__version__, tensordict.__version__, torchrl.__version__, numpy.__version__, sys.version, sys.platform)

2.2.0a0+81ea7a4 0.4.0+eaef29e 0.4.0+01a2216 1.24.4 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] linux
```


Discussion:
If you change the env to `SerialEnv` or remove the `TransformedEnv` on the ParallelEnv, the error goes away. This is a bit strange.
If you change the truncated key to something else than `truncated` the error goes away too. This makes sense since the original truncated key could be written natively by the TimeLimitWrapper on the CPU env. But why would the local StepCounter write on the BatchedEnv device?

Resolving patch:
diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
--- a/torchrl/envs/batched_envs.py
+++ b/torchrl/envs/batched_envs.py
@@ -1284,7 +1284,22 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
                         tensordict_, keys_to_update=list(self._selected_reset_keys)
                     )
                 continue
-            out = ("reset", tensordict_)
+            if tensordict_ is not None:
+                tdkeys = list(tensordict_.keys(True, True))
+
+                # This way we can avoid calling select over all the keys in the shared tensordict
+                def tentative_update(val, other):
+                    if other is not None:
+                        val.copy_(other)
+                    return val
+
+                self.shared_tensordicts[i].apply_(
+                    tentative_update, tensordict_, default=None
+                )
+                out = ("reset", tdkeys)
+            else:
+                out = ("reset", False)
+
             channel.send(out)
             workers.append(i)
 
@@ -1509,7 +1524,13 @@ def look_for_cuda(tensor, has_cuda=has_cuda):
                 torchrl_logger.info(f"resetting worker {pid}")
             if not initialized:
                 raise RuntimeError("call 'init' before resetting")
-            cur_td = env.reset(tensordict=data)
+            # we use 'data' to pass the keys that we need to pass to reset,
+            # because passing the entire buffer may have unwanted consequences
+            cur_td = env.reset(
+                tensordict=root_shared_tensordict.select(*data, strict=False)
+                if data
+                else None
+            )
             shared_tensordict.update_(
                 cur_td,
                 keys_to_update=list(_selected_reset_keys),


Test patch:
diff --git a/test/test_env.py b/test/test_env.py
--- a/test/test_env.py
+++ b/test/test_env.py
@@ -624,6 +624,29 @@ def test_parallel_env_with_policy(
         # env_serial.close()
         env0.close()
 
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+    @pytest.mark.parametrize("heterogeneous", [False, True])
+    def test_transform_env_transform_no_device(self, heterogeneous):
+        # Tests non-regression on 1865
+        def make_env():
+            return TransformedEnv(
+                ContinuousActionVecMockEnv(), StepCounter(max_steps=3)
+            )
+
+        if heterogeneous:
+            make_envs = [EnvCreator(make_env), EnvCreator(make_env)]
+        else:
+            make_envs = make_env
+        penv = ParallelEnv(2, make_envs)
+        r = penv.rollout(6, break_when_any_done=False)
+        assert r.shape == (2, 6)
+        try:
+            env = TransformedEnv(penv)
+            r = env.rollout(6, break_when_any_done=False)
+            assert r.shape == (2, 6)
+        finally:
+            penv.close()
+
     @pytest.mark.skipif(not _has_gym, reason="no gym")
     @pytest.mark.parametrize(
         "env_name",