[BugFix] Sync with tensordict (meta-tensor deprecation) (#842)

Author: vmoens

test/test_shared.py

@@ -9,7 +9,7 @@
 
 import pytest
 import torch
-from tensordict import SavedTensorDict, TensorDict
+from tensordict import TensorDict
 from torch import multiprocessing as mp
 
 
@@ -145,10 +145,10 @@ def test_shared(self, shared):
             )
 
 
-@pytest.mark.skipif(
-    sys.platform == "win32",
-    reason="RuntimeError from Torch serialization.py when creating td_saved on Windows",
-)
+# @pytest.mark.skipif(
+#     sys.platform == "win32",
+#     reason="RuntimeError from Torch serialization.py when creating td_saved on Windows",
+# )
 @pytest.mark.parametrize(
     "idx",
     [
@@ -180,7 +180,6 @@ def test_memmap(idx, dtype, large_scale=False):
 
     td_sm = td.clone().share_memory_()
     td_memmap = td.clone().memmap_()
-    td_saved = td.to(SavedTensorDict)
 
     print("\nTesting reading from TD")
     for i in range(2):
@@ -194,11 +193,6 @@ def test_memmap(idx, dtype, large_scale=False):
         if i == 1:
             print(f"memmap: {time.time() - t0:4.4f} sec")
 
-        t0 = time.time()
-        td_saved[idx].clone()
-        if i == 1:
-            print(f"saved td: {time.time() - t0:4.4f} sec")
-
     td_to_copy = td[idx].contiguous()
     for k in td_to_copy.keys():
         td_to_copy.set_(k, torch.ones_like(td_to_copy.get(k)))
@@ -219,13 +213,6 @@ def test_memmap(idx, dtype, large_scale=False):
             print(f"memmap td: {time.time() - t0:4.4f} sec")
         torch.testing.assert_close(sub_td_sm.get("a")._tensor, td_to_copy.get("a"))
 
-        t0 = time.time()
-        sub_td_sm = td_saved.get_sub_tensordict(idx)
-        sub_td_sm.update_(td_to_copy)
-        if i == 1:
-            print(f"saved td: {time.time() - t0:4.4f} sec")
-        torch.testing.assert_close(sub_td_sm.get("a"), td_to_copy.get("a"))
-
 
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()

torchrl/collectors/utils.py

@@ -59,7 +59,7 @@ def split_trajectories(rollout_tensordict: TensorDictBase) -> TensorDictBase:
             torch.ones(
                 out_split.shape,
                 dtype=torch.bool,
-                device=out_split._get_meta("done").device,
+                device=out_split.get("done").device,
             ),
         )
     MAX = max(*[out_split.shape[0] for out_split in out_splits])

torchrl/envs/common.py

@@ -380,13 +380,13 @@ def step(self, tensordict: TensorDictBase) -> TensorDictBase:
                 obs = tensordict_out.get(key)
                 self.observation_spec.type_check(obs, key)
 
-            if tensordict_out._get_meta("reward").dtype is not self.reward_spec.dtype:
+            if tensordict_out.get("reward").dtype is not self.reward_spec.dtype:
                 raise TypeError(
                     f"expected reward.dtype to be {self.reward_spec.dtype} "
                     f"but got {tensordict_out.get('reward').dtype}"
                 )
 
-            if tensordict_out._get_meta("done").dtype is not torch.bool:
+            if tensordict_out.get("done").dtype is not torch.bool:
                 raise TypeError(
                     f"expected done.dtype to be torch.bool but got {tensordict_out.get('done').dtype}"
                 )