[BugFix] Fix MultOneHotDiscreteTensorSpec.is_in (#818)

Author: riiswa

test/test_tensor_spec.py

@@ -231,6 +231,7 @@ def test_mult_onehot(shape, ns):
             assert (_r.sum(-1) == 1).all()
             assert _r.shape[-1] == _n
         np_r = ts.to_numpy(r)
+        assert not ts.is_in(torch.tensor(np_r))
         assert (ts.encode(np_r) == r).all()
 
 

torchrl/data/tensor_specs.py

@@ -878,9 +878,11 @@ def encode(self, val: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
             x.append(super(MultOneHotDiscreteTensorSpec, self).encode(v, space))
         return torch.cat(x, -1)
 
-    def _split(self, val: torch.Tensor) -> torch.Tensor:
-        vals = val.split([space.n for space in self.space], dim=-1)
-        return vals
+    def _split(self, val: torch.Tensor) -> Optional[torch.Tensor]:
+        split_sizes = [space.n for space in self.space]
+        if val.ndim < 1 or val.shape[-1] != sum(split_sizes):
+            return None
+        return val.split(split_sizes, dim=-1)
 
     def to_numpy(self, val: torch.Tensor, safe: bool = True) -> np.ndarray:
         if safe:
@@ -907,8 +909,10 @@ def index(self, index: INDEX_TYPING, tensor_to_index: torch.Tensor) -> torch.Ten
 
     def is_in(self, val: torch.Tensor) -> bool:
         vals = self._split(val)
+        if vals is None:
+            return False
         return all(
-            [super(MultOneHotDiscreteTensorSpec, self).is_in(_val) for _val in vals]
+            super(MultOneHotDiscreteTensorSpec, self).is_in(_val) for _val in vals
         )
 
     def _project(self, val: torch.Tensor) -> torch.Tensor: