[Refactor] LAZY_LEGACY_OP=False (#1832)

Vincent Moens · web-flow · commit 9da61f239258 · 2024-01-29T09:12:10.000Z
diff --git a/.github/unittest/linux/scripts/run_all.sh b/.github/unittest/linux/scripts/run_all.sh
@@ -76,7 +76,7 @@ export DISPLAY=:0
 export SDL_VIDEODRIVER=dummy
 
 # legacy from bash scripts: remove?
-conda env config vars set MUJOCO_GL=$MUJOCO_GL PYOPENGL_PLATFORM=$MUJOCO_GL DISPLAY=:0 SDL_VIDEODRIVER=dummy
+conda env config vars set MUJOCO_GL=$MUJOCO_GL PYOPENGL_PLATFORM=$MUJOCO_GL DISPLAY=:0 SDL_VIDEODRIVER=dummy LAZY_LEGACY_OP=False
 
 pip3 install pip --upgrade
 pip install virtualenv
diff --git a/.github/unittest/linux_distributed/scripts/run_test.sh b/.github/unittest/linux_distributed/scripts/run_test.sh
@@ -18,6 +18,7 @@ lib_dir="${env_dir}/lib"
 export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$lib_dir
 export MKL_THREADING_LAYER=GNU
 export CKPT_BACKEND=torch
+export LAZY_LEGACY_OP=False
 export BATCHED_PIPE_TIMEOUT=60
 
 python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/smoke_test.py -v --durations 200
diff --git a/.github/unittest/linux_examples/scripts/run_all.sh b/.github/unittest/linux_examples/scripts/run_all.sh
@@ -74,6 +74,7 @@ export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$root_dir/.mujoco/mujoco210/bin
 export SDL_VIDEODRIVER=dummy
 export MUJOCO_GL=egl
 export PYOPENGL_PLATFORM=egl
+export LAZY_LEGACY_OP=False
 
 conda env config vars set MUJOCO_PY_MUJOCO_PATH=$root_dir/.mujoco/mujoco210 \
   DISPLAY=unix:0.0 \
diff --git a/.github/unittest/linux_libs/scripts_ataridqn/run_test.sh b/.github/unittest/linux_libs/scripts_ataridqn/run_test.sh
@@ -8,6 +8,7 @@ conda activate ./env
 apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf libosmesa6-dev libgl1-mesa-glx libglfw3 swig3.0
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
+export LAZY_LEGACY_OP=False
 export PYTORCH_TEST_WITH_SLOW='1'
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
diff --git a/.github/unittest/linux_libs/scripts_brax/run_test.sh b/.github/unittest/linux_libs/scripts_brax/run_test.sh
@@ -7,6 +7,7 @@ conda activate ./env
 
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_d4rl/run_test.sh b/.github/unittest/linux_libs/scripts_d4rl/run_test.sh
@@ -24,6 +24,7 @@ cd ..
 #cd ..
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_envpool/run_test.sh b/.github/unittest/linux_libs/scripts_envpool/run_test.sh
@@ -12,6 +12,7 @@ eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_gen-dgrl/run_test.sh b/.github/unittest/linux_libs/scripts_gen-dgrl/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_gym/run_test.sh b/.github/unittest/linux_libs/scripts_gym/run_test.sh
@@ -6,6 +6,7 @@ eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_habitat/run_test.sh b/.github/unittest/linux_libs/scripts_habitat/run_test.sh
@@ -16,6 +16,7 @@ conda env config vars set LD_PRELOAD=$LD_PRELOAD:$STDC_LOC
 STDC_LOC=$(find conda/ -name "libstdc++.so.6" | head -1)
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_jumanji/run_test.sh b/.github/unittest/linux_libs/scripts_jumanji/run_test.sh
@@ -8,6 +8,7 @@ apt-get update && apt-get install -y git wget
 
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_minari/run_test.sh b/.github/unittest/linux_libs/scripts_minari/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_openx/run_test.sh b/.github/unittest/linux_libs/scripts_openx/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_pettingzoo/run_test.sh b/.github/unittest/linux_libs/scripts_pettingzoo/run_test.sh
@@ -8,6 +8,7 @@ apt-get update && apt-get install -y git wget
 
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_rlhf/run_test.sh b/.github/unittest/linux_libs/scripts_rlhf/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get install -y git gcc
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_robohive/setup_env.sh b/.github/unittest/linux_libs/scripts_robohive/setup_env.sh
@@ -66,7 +66,8 @@ conda env config vars set \
   DISPLAY=unix:0.0 \
   PYOPENGL_PLATFORM=egl \
   NVIDIA_PATH=/usr/src/nvidia-470.63.01 \
-  sim_backend=MUJOCO
+  sim_backend=MUJOCO \
+  LAZY_LEGACY_OP=False
 
 # make env variables apparent
 conda deactivate && conda activate "${env_dir}"
diff --git a/.github/unittest/linux_libs/scripts_roboset/run_test.sh b/.github/unittest/linux_libs/scripts_roboset/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_sklearn/run_test.sh b/.github/unittest/linux_libs/scripts_sklearn/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get install -y git gcc
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_smacv2/run_test.sh b/.github/unittest/linux_libs/scripts_smacv2/run_test.sh
@@ -8,6 +8,7 @@ apt-get update && apt-get install -y git wget
 
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_vd4rl/run_test.sh b/.github/unittest/linux_libs/scripts_vd4rl/run_test.sh
@@ -9,6 +9,7 @@ apt-get update && apt-get remove swig -y && apt-get install -y git gcc patchelf
 ln -s /usr/bin/swig3.0 /usr/bin/swig
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_libs/scripts_vmas/run_test.sh b/.github/unittest/linux_libs/scripts_vmas/run_test.sh
@@ -8,6 +8,7 @@ apt-get update && apt-get install -y git wget
 
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_olddeps/scripts_gym_0_13/run_test.sh b/.github/unittest/linux_olddeps/scripts_gym_0_13/run_test.sh
@@ -6,6 +6,7 @@ eval "$(./conda/bin/conda shell.bash hook)"
 conda activate ./env
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/linux_optdeps/scripts/run_test.sh b/.github/unittest/linux_optdeps/scripts/run_test.sh
@@ -9,6 +9,7 @@ conda activate ./env
 STDC_LOC=$(find conda/ -name "libstdc++.so.6" | head -1)
 
 export PYTORCH_TEST_WITH_SLOW='1'
+export LAZY_LEGACY_OP=False
 python -m torch.utils.collect_env
 # Avoid error: "fatal: unsafe repository"
 git config --global --add safe.directory '*'
diff --git a/.github/unittest/windows_optdepts/scripts/run_test.sh b/.github/unittest/windows_optdepts/scripts/run_test.sh
@@ -12,6 +12,7 @@ source "$this_dir/set_cuda_envs.sh"
 export CKPT_BACKEND=torch
 export MAX_IDLE_COUNT=60
 export BATCHED_PIPE_TIMEOUT=60
+export LAZY_LEGACY_OP=False
 
 python -m torch.utils.collect_env
 pytest --junitxml=test-results/junit.xml -v --durations 200  --ignore test/test_distributed.py --ignore test/test_rlhf.py
diff --git a/test/test_collector.py b/test/test_collector.py
@@ -35,6 +35,7 @@
     MultiKeyCountingEnvPolicy,
     NestedCountingEnv,
 )
+from tensordict import LazyStackedTensorDict
 from tensordict.nn import TensorDictModule, TensorDictSequential
 from tensordict.tensordict import assert_allclose_td, TensorDict
 
@@ -1896,7 +1897,7 @@ def test_aggregate_reset_to_root(self):
             },
             [1, 2],
         )
-        td = torch.stack([td0, td1], 0)
+        td = LazyStackedTensorDict.lazy_stack([td0, td1], 0)
         assert _aggregate_end_of_traj(td).all()
 
     def test_aggregate_reset_to_root_keys(self):
@@ -1991,7 +1992,7 @@ def test_aggregate_reset_to_root_keys(self):
             },
             [1, 2],
         )
-        td = torch.stack([td0, td1], 0)
+        td = LazyStackedTensorDict.lazy_stack([td0, td1], 0)
         assert _aggregate_end_of_traj(td, reset_keys=["_reset", ("a", "_reset")]).all()
 
     def test_aggregate_reset_to_root_errors(self):
diff --git a/test/test_env.py b/test/test_env.py
@@ -1141,10 +1141,12 @@ def test_steptensordict(
             tds[0]["this", "one"] = torch.zeros(2)
             tds[1]["but", "not", "this", "one"] = torch.ones(2)
             tds[0]["next", "this", "one"] = torch.ones(2) * 2
-            tensordict = torch.stack(tds, 0)
+            tensordict = LazyStackedTensorDict.lazy_stack(tds, 0)
         next_tensordict = TensorDict({}, [4]) if has_out else None
         if has_out and lazy_stack:
-            next_tensordict = torch.stack(next_tensordict.unbind(0), 0)
+            next_tensordict = LazyStackedTensorDict.lazy_stack(
+                next_tensordict.unbind(0), 0
+            )
         out = step_mdp(
             tensordict.lock_(),
             keep_other=keep_other,
@@ -1498,8 +1500,7 @@ def test_heterogeenous(
                     [td_batch_size],
                 )
             )
-        lazy_td = torch.stack(tds, dim=1)
-        input_td = lazy_td
+        lazy_td = LazyStackedTensorDict.lazy_stack(tds, dim=1)
 
         td = step_mdp(
             lazy_td.lock_(),
@@ -1785,7 +1786,7 @@ def main_penv(j, q=None):
                 r_p.append(env_s.rollout(100, break_when_any_done=False, policy=policy))
                 r_s.append(env_p.rollout(100, break_when_any_done=False, policy=policy))
 
-        td_equals = torch.stack(r_p).contiguous() == torch.stack(r_s).contiguous()
+        td_equals = torch.stack(r_p) == torch.stack(r_s)
         if td_equals.all():
             if q is not None:
                 q.put(("passed", j))
diff --git a/test/test_libs.py b/test/test_libs.py
@@ -555,7 +555,7 @@ def non_null_obs(batched_td):
             env_type = type(env0._env)
 
         assert_allclose_td(*tdreset, rtol=RTOL, atol=ATOL)
-        tdrollout = torch.stack(tdrollout, 0).contiguous()
+        tdrollout = torch.stack(tdrollout, 0)
 
         # custom filtering of non-null obs: mujoco rendering sometimes fails
         # and renders black images. To counter this in the tests, we select
@@ -597,7 +597,7 @@ def non_null_obs(batched_td):
         assert_allclose_td(tdreset[0], tdreset2, rtol=RTOL, atol=ATOL)
         assert final_seed0 == final_seed2
         # same magic trick for mujoco as above
-        tdrollout = torch.stack([tdrollout[0], rollout2], 0).contiguous()
+        tdrollout = torch.stack([tdrollout[0], rollout2], 0)
         idx = non_null_obs(tdrollout)
         assert_allclose_td(
             tdrollout[0][..., idx], tdrollout[1][..., idx], rtol=RTOL, atol=ATOL
diff --git a/test/test_shared.py b/test/test_shared.py
@@ -9,7 +9,7 @@
 
 import pytest
 import torch
-from tensordict import TensorDict
+from tensordict import LazyStackedTensorDict, TensorDict
 from torch import multiprocessing as mp
 
 
@@ -81,7 +81,7 @@ def remote_process(command_pipe_child, command_pipe_parent, tensordict):
         command_pipe_parent.close()
         assert isinstance(tensordict, TensorDict), f"td is of type {type(tensordict)}"
         assert tensordict.is_shared() or tensordict.is_memmap()
-        new_tensordict = torch.stack(
+        new_tensordict = LazyStackedTensorDict.lazy_stack(
             [
                 tensordict[i].contiguous().clone().zero_()
                 for i in range(tensordict.shape[0])
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -791,7 +791,7 @@ def test_transform_model(self, dim, N, padding):
             model(tdbase0)
         tdbase0.batch_size = [10]
         tdbase0 = tdbase0.expand(5, 10)
-        tdbase0_copy = tdbase0.transpose(0, 1).to_tensordict()
+        tdbase0_copy = tdbase0.transpose(0, 1)
         tdbase0.refine_names("time", None)
         tdbase0_copy.names = [None, "time"]
         v1 = model(tdbase0)
diff --git a/torchrl/__init__.py b/torchrl/__init__.py
@@ -11,7 +11,7 @@
 
 from torch import multiprocessing as mp
 
-set_lazy_legacy(True).set()
+set_lazy_legacy(False).set()
 
 if torch.cuda.device_count() > 1:
     n = torch.cuda.device_count() - 1
diff --git a/torchrl/collectors/utils.py b/torchrl/collectors/utils.py
@@ -6,6 +6,8 @@
 from typing import Callable
 
 import torch
+
+from tensordict import set_lazy_legacy
 from tensordict.tensordict import pad, TensorDictBase
 
 
@@ -25,6 +27,7 @@ def stacked_output_fun(*args, **kwargs):
     return stacked_output_fun
 
 
+@set_lazy_legacy(False)
 def split_trajectories(
     rollout_tensordict: TensorDictBase, prefix=None
 ) -> TensorDictBase:
@@ -88,7 +91,7 @@ def split_trajectories(
             ),
         )
         if rollout_tensordict.ndimension() == 1:
-            rollout_tensordict = rollout_tensordict.unsqueeze(0).to_tensordict()
+            rollout_tensordict = rollout_tensordict.unsqueeze(0)
         return rollout_tensordict.unflatten_keys(sep)
     out_splits = rollout_tensordict.view(-1).split(splits, 0)
 
diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -348,7 +348,7 @@ def _check_for_empty_spec(specs: CompositeSpec):
             self.done_spec = output_spec["full_done_spec"]
 
             self._dummy_env_str = str(meta_data[0])
-            self._env_tensordict = torch.stack(
+            self._env_tensordict = LazyStackedTensorDict.lazy_stack(
                 [meta_data.tensordict for meta_data in meta_data], 0
             )
             self._batch_locked = meta_data[0].batch_locked
@@ -463,7 +463,7 @@ def _create_td(self) -> None:
                 )
                 for tensordict in shared_tensordict_parent
             ]
-            shared_tensordict_parent = torch.stack(
+            shared_tensordict_parent = LazyStackedTensorDict.lazy_stack(
                 shared_tensordict_parent,
                 0,
             )
@@ -474,7 +474,9 @@ def _create_td(self) -> None:
                 self.shared_tensordicts = [
                     td.clone() for td in self.shared_tensordict_parent.unbind(0)
                 ]
-                self.shared_tensordict_parent = torch.stack(self.shared_tensordicts, 0)
+                self.shared_tensordict_parent = LazyStackedTensorDict.lazy_stack(
+                    self.shared_tensordicts, 0
+                )
             else:
                 # Multi-task: we share tensordict that *may* have different keys
                 # LazyStacked already stores this so we don't need to do anything
diff --git a/torchrl/envs/common.py b/torchrl/envs/common.py
@@ -14,7 +14,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from tensordict import unravel_key
+from tensordict import LazyStackedTensorDict, unravel_key
 from tensordict.tensordict import TensorDictBase
 from tensordict.utils import NestedKey
 from torchrl._utils import _replace_last, implement_for, prod, seed_generator
@@ -2307,9 +2307,12 @@ def rollout(
         else:
             tensordicts = self._rollout_nonstop(**kwargs)
         batch_size = self.batch_size if tensordict is None else tensordict.batch_size
-        out_td = torch.stack(tensordicts, len(batch_size), out=out)
         if return_contiguous:
-            out_td = out_td.contiguous()
+            out_td = torch.stack(tensordicts, len(batch_size), out=out)
+        else:
+            out_td = LazyStackedTensorDict.lazy_stack(
+                tensordicts, len(batch_size), out=out
+            )
         out_td.refine_names(..., "time")
         return out_td
 
@@ -2408,7 +2411,7 @@ def step_and_maybe_reset(
             ...     for i in range(n):
             ...         data, data_ = env.step_and_maybe_reset(data_)
             ...         result.append(data)
-            ...     return torch.stack(result).contiguous()
+            ...     return torch.stack(result)
             >>> env = ParallelEnv(2, lambda: GymEnv("CartPole-v1"))
             >>> print(rollout(env, 2))
             TensorDict(
diff --git a/torchrl/envs/transforms/r3m.py b/torchrl/envs/transforms/r3m.py
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
diff --git a/torchrl/envs/transforms/vc1.py b/torchrl/envs/transforms/vc1.py
diff --git a/torchrl/envs/transforms/vip.py b/torchrl/envs/transforms/vip.py
diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py
diff --git a/torchrl/trainers/helpers/models.py b/torchrl/trainers/helpers/models.py
