[BugFix] Preserve action masks when expanding categorical specs

test/test_specs.py

@@ -1760,6 +1760,18 @@ def test_discrete(self, shape1, shape2):
         assert spec2.rand().shape == spec2.shape
         assert spec2.zero().shape == spec2.shape
 
+    @pytest.mark.parametrize("batch_shape", [(2, 4), (3, 2, 4)])
+    def test_discrete_masked(self, batch_shape):
+        mask = torch.tensor([[False, True, False], [False, False, True]]).unsqueeze(1)
+        spec = Categorical(3, shape=(2, 1), mask=mask)
+
+        expanded = spec.expand(batch_shape)
+        torch.testing.assert_close(expanded.mask, mask.expand(*batch_shape, 3))
+        sample = expanded.rand((5,))
+        expected = torch.tensor([[1], [2]]).expand(5, *batch_shape)
+        torch.testing.assert_close(sample, expected)
+        assert spec.is_in(sample)
+
     @pytest.mark.parametrize("shape1", [None, (), (5,)])
     @pytest.mark.parametrize("shape2", [(), (10,)])
     def test_multidiscrete(self, shape1, shape2):
@@ -1788,6 +1800,20 @@ def test_multidiscrete(self, shape1, shape2):
         assert spec2.rand().shape == spec2.shape
         assert spec2.zero().shape == spec2.shape
 
+    @pytest.mark.parametrize("batch_shape", [(2, 4), (3, 2, 4)])
+    def test_multidiscrete_masked(self, batch_shape):
+        mask = torch.tensor(
+            [[False, True, False, True, False], [False, False, True, False, True]]
+        ).unsqueeze(1)
+        spec = MultiCategorical([3, 2], shape=(2, 1, 2), mask=mask)
+
+        expanded = spec.expand(*batch_shape, 2)
+        torch.testing.assert_close(expanded.mask, mask.expand(*batch_shape, 5))
+        sample = expanded.rand((5,))
+        expected = torch.tensor([[[1, 0]], [[2, 1]]]).expand(5, *batch_shape, 2)
+        torch.testing.assert_close(sample, expected)
+        assert spec.is_in(sample)
+
     @pytest.mark.parametrize("shape1", [None, (), (5,)])
     @pytest.mark.parametrize("shape2", [(), (10,)])
     def test_multionehot(self, shape1, shape2):

torchrl/data/tensor_specs.py

@@ -955,6 +955,8 @@ def expand(self, shape: torch.Size):
     def expand(self, *shape: int) -> T:
         """Returns a new Spec with the expanded shape.
 
+        If the spec has an outcome mask, it is broadcast to the expanded shape.
+
         Args:
             *shape (tuple or iterable of int): the new shape of the Spec.
                 Must be broadcastable with the current shape:
@@ -4174,7 +4176,11 @@ def expand(self, *shape):
                 f"shape of the {self.__class__.__name__} spec in expand()."
             )
         return self.__class__(
-            n=self.space.n, shape=shape, device=self.device, dtype=self.dtype
+            n=self.space.n,
+            shape=shape,
+            device=self.device,
+            dtype=self.dtype,
+            mask=self.mask,
         )
 
     def _reshape(self, shape):
@@ -4983,17 +4989,12 @@ def expand(self, *shape):
                 f"The last {self.ndim} of the expanded shape {shape} must match the"
                 f"shape of the {self.__class__.__name__} spec in expand()."
             )
-        mask = (
-            self.mask.expand(*shape, self.mask.shape[-1])
-            if self.mask is not None
-            else None
-        )
         return self.__class__(
             nvec=self.nvec,
             shape=shape,
             device=self.device,
             dtype=self.dtype,
-            mask=mask,
+            mask=self.mask,
         )
 
     def _reshape(self, shape):