amend

pytorch · vmoens · May 31, 2024 · Apr 30, 2024 · May 26, 2024 · May 26, 2024
commit b0c208f8ae89243c13b4ec83189970c8a5216b8f
diff --git a/test/test_libs.py b/test/test_libs.py
@@ -365,9 +365,10 @@ def stack_tuples(item):
         # sample_pt = torch.utils._pytree.tree_map(lambda x: torch.stack(list(x)), sample_pt, is_leaf=lambda x: isinstance(x, tuple))
         spec = _gym_to_torchrl_spec_transform(space)
         rand = spec.rand()
-        assert spec.is_in(rand)
-        assert rand in spec
-        assert sample_pt in spec, (rand, sample_pt)
+
+        assert spec.contains(rand), (rand, spec)
+        assert spec.contains(sample_pt), (rand, sample_pt)
+
         space_recon = _torchrl_to_gym_spec_transform(spec)
         assert space_recon == space, (space_recon, space)
         rand_numpy = rand.numpy()

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -1254,8 +1254,15 @@ def type_check(self, value: torch.Tensor, key: str = None) -> None:
             spec.type_check(val)
 
     def is_in(self, value) -> bool:
-        # We don't use unbind because value could be a tuple
-        return all(value in spec for (value, spec) in zip(value, self._specs))
+        if self.dim == 0:
+            # We don't use unbind because value could be a tuple or a nested tensor
+            return all(
+                spec.contains(value) for (value, spec) in zip(value, self._specs)
+            )
+        return all(
+            spec.contains(value)
+            for (value, spec) in zip(value.unbind(self.dim), self._specs)
+        )
 
     @property
     def space(self):