Repository: pytorch/rl

Issue:
[BUG] MultOneHotDiscreteTensorSpec.is_in raise a RuntimeError when input is not a one hot tensor 
## Describe the bug

MultOneHotDiscreteTensorSpec.is_in raise a RuntimeError when input is not a one hot tensor.

```python
# The input is tensor([[0, 1, 0],
#        [0, 0, 1],
#      [3, 0, 1]])
test/test_tensor_spec.py:310 (test_multi_discrete_conversion[device0-shape0-ns1])
Traceback (most recent call last):
  File "/Users/waris/Projects/rl/test/test_tensor_spec.py", line 338, in test_multi_discrete_conversion
    assert one_hot.is_in(categorical.to_onehot(categorical.rand(shape)))
  File "/Users/waris/Projects/rl/torchrl/data/tensor_specs.py", line 909, in is_in
    vals = self._split(val)
  File "/Users/waris/Projects/rl/torchrl/data/tensor_specs.py", line 882, in _split
    vals = val.split([space.n for space in self.space], dim=-1)
  File "/Users/waris/Projects/rl/venv/lib/python3.10/site-packages/torch/_tensor.py", line 789, in split
    return torch._VF.split_with_sizes(self, split_size, dim)
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 3 (input tensor's size at dimension -1), but got split_sizes=[5, 2, 3]
```

## Expected behavior

Should return `False` instead raising an error. 

**NOTE**: In general when we test method that return boolean, we should also write tests for the `False` case.

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
Good catch
Is there a clever way of doing this or should we just capture the error and return False in that case?
> NOTE: In general when we test method that return boolean, we should also write tests for the False case.

hehe, indeed :) 
> Good catch Is there a clever way of doing this or should we just capture the error and return False in that case?

I'm gonna work on a simple and clever way of doing it, because exception handling can lead to performance issue.
+ undefined behaviour and/or silent errors

Resolving patch:
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
--- a/torchrl/data/tensor_specs.py
+++ b/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:


Test patch:
diff --git a/test/test_tensor_spec.py b/test/test_tensor_spec.py
--- a/test/test_tensor_spec.py
+++ b/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()