Update

[ghstack-poisoned]

Author: kurtamohler

test/test_specs.py

@@ -689,17 +689,56 @@ class TestChoiceSpec:
     @pytest.mark.parametrize("input_type", ["spec", "nontensor", "nontensorstack"])
     def test_choice(self, input_type):
         if input_type == "spec":
-            stack = torch.stack([Bounded(0, 2.5, ()), Bounded(10, 12, ())])
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
+            example_in = torch.tensor(11.0)
+            example_out = torch.tensor(9.0)
         elif input_type == "nontensor":
-            stack = torch.stack([NonTensorData("a"), NonTensorData("b")])
+            choices = [NonTensorData("a"), NonTensorData("b")]
+            example_in = NonTensorData("b")
+            example_out = NonTensorData("c")
         elif input_type == "nontensorstack":
-            stack = torch.stack(
-                [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
-            )
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
+            example_in = NonTensorStack("a", "b", "c")
+            example_out = NonTensorStack("a", "c", "b")
 
-        spec = Choice(stack)
+        spec = Choice(choices)
         res = spec.rand()
         assert spec.is_in(res)
+        assert spec.is_in(example_in)
+        assert not spec.is_in(example_out)
+
+    def test_errors(self):
+        with pytest.raises(TypeError, match="must be a list"):
+            Choice("abc")
+
+        with pytest.raises(
+            TypeError,
+            match="must be either a TensorSpec, NonTensorData, or NonTensorStack",
+        ):
+            Choice(["abc"])
+
+        with pytest.raises(TypeError, match="must be the same type"):
+            Choice([Bounded(0, 1, (1,)), Categorical(10, (1,))])
+
+        with pytest.raises(ValueError, match="must have the same shape"):
+            Choice([Categorical(10, (1,)), Categorical(10, (2,))])
+
+        with pytest.raises(ValueError, match="must have the same dtype"):
+            Choice(
+                [
+                    Categorical(10, (2,), dtype=torch.long),
+                    Categorical(10, (2,), dtype=torch.float),
+                ]
+            )
+
+        if torch.cuda.device_count():
+            with pytest.raises(ValueError, match="must have the same device"):
+                Choice(
+                    [
+                        Categorical(10, (2,), device="cpu"),
+                        Categorical(10, (2,), device="cuda"),
+                    ]
+                )
 
 
 @pytest.mark.parametrize("shape", [(), (2, 3)])
@@ -1436,17 +1475,23 @@ def test_non_tensor(self):
     @pytest.mark.parametrize("input_type", ["spec", "nontensor", "nontensorstack"])
     def test_choice(self, input_type):
         if input_type == "spec":
-            stack = torch.stack([Bounded(0, 2.5, ()), Bounded(10, 12, ())])
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
         elif input_type == "nontensor":
-            stack = torch.stack([NonTensorData("a"), NonTensorData("b")])
+            choices = [NonTensorData("a"), NonTensorData("b")]
         elif input_type == "nontensorstack":
-            stack = torch.stack(
-                [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
-            )
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
 
-        spec = Choice(stack)
-        with pytest.raises(NotImplementedError):
-            spec.expand((3,))
+        spec = Choice(choices)
+        res = spec.expand(
+            [
+                3,
+            ]
+        )
+        assert res.shape == torch.Size(
+            [
+                3,
+            ]
+        )
 
     @pytest.mark.parametrize("shape1", [None, (), (5,)])
     @pytest.mark.parametrize("shape2", [(), (10,)])
@@ -1653,15 +1698,13 @@ def test_non_tensor(self):
     @pytest.mark.parametrize("input_type", ["spec", "nontensor", "nontensorstack"])
     def test_choice(self, input_type):
         if input_type == "spec":
-            stack = torch.stack([Bounded(0, 2.5, ()), Bounded(10, 12, ())])
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
         elif input_type == "nontensor":
-            stack = torch.stack([NonTensorData("a"), NonTensorData("b")])
+            choices = [NonTensorData("a"), NonTensorData("b")]
         elif input_type == "nontensorstack":
-            stack = torch.stack(
-                [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
-            )
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
 
-        spec = Choice(stack)
+        spec = Choice(choices)
         assert spec.clone() == spec
         assert spec.clone() is not spec
 
@@ -1756,19 +1799,15 @@ def test_non_tensor(self):
     )
     def test_choice(self, input_type):
         if input_type == "bounded_spec":
-            stack = torch.stack([Bounded(0, 2.5, ()), Bounded(10, 12, ())])
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
         elif input_type == "categorical_spec":
-            stack = torch.stack([Categorical(10), Categorical(20)])
+            choices = [Categorical(10), Categorical(20)]
         elif input_type == "nontensor":
-            stack = torch.stack(
-                [NonTensorData("a"), NonTensorData("b"), NonTensorData("c")]
-            )
+            choices = [NonTensorData("a"), NonTensorData("b"), NonTensorData("c")]
         elif input_type == "nontensorstack":
-            stack = torch.stack(
-                [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
-            )
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
 
-        spec = Choice(stack)
+        spec = Choice(choices)
 
         if input_type == "bounded_spec":
             assert spec.cardinality() == float("inf")
@@ -2093,19 +2132,15 @@ def test_non_tensor(self, device):
     )
     def test_choice(self, input_type, device):
         if input_type == "bounded_spec":
-            stack = torch.stack([Bounded(0, 2.5, ()), Bounded(10, 12, ())])
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
         elif input_type == "categorical_spec":
-            stack = torch.stack([Categorical(10), Categorical(20)])
+            choices = [Categorical(10), Categorical(20)]
         elif input_type == "nontensor":
-            stack = torch.stack(
-                [NonTensorData("a"), NonTensorData("b"), NonTensorData("c")]
-            )
+            choices = [NonTensorData("a"), NonTensorData("b")]
         elif input_type == "nontensorstack":
-            stack = torch.stack(
-                [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
-            )
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
 
-        spec = Choice(stack, device="cpu")
+        spec = Choice(choices)
         assert spec.to(device).device == device
 
     @pytest.mark.parametrize("shape1", [(5,), (5, 6)])
@@ -2376,26 +2411,30 @@ def test_stack_non_tensor(self, shape, stack_dim):
 
     @pytest.mark.parametrize(
         "input_type",
-        ["bounded_spec", "categorical_spec", "nontensor", "nontensorstack"],
+        ["bounded_spec", "categorical_spec", "nontensor"],
     )
     def test_stack_choice(self, input_type, shape, stack_dim):
         if input_type == "bounded_spec":
-            stack = torch.stack([Bounded(0, 2.5, ()), Bounded(10, 12, ())])
+            choices = [Bounded(0, 2.5, shape), Bounded(10, 12, shape)]
         elif input_type == "categorical_spec":
-            stack = torch.stack([Categorical(10), Categorical(20)])
+            choices = [Categorical(10, shape), Categorical(20, shape)]
         elif input_type == "nontensor":
-            stack = torch.stack(
-                [NonTensorData("a"), NonTensorData("b"), NonTensorData("c")]
-            )
-        elif input_type == "nontensorstack":
-            stack = torch.stack(
-                [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
-            )
+            if len(shape) == 0:
+                choices = [NonTensorData("a"), NonTensorData("b"), NonTensorData("c")]
+            else:
+                choices = [
+                    NonTensorStack("a").expand(shape + (1,)).squeeze(-1),
+                    NonTensorStack("d").expand(shape + (1,)).squeeze(-1),
+                ]
 
-        spec0 = Choice(stack)
-        spec1 = Choice(stack)
-        with pytest.raises(NotImplementedError):
-            torch.stack([spec0, spec1], 0)
+        spec0 = Choice(choices)
+        spec1 = Choice(choices)
+        res = torch.stack([spec0, spec1], stack_dim)
+        assert isinstance(res, Choice)
+        assert (
+            res.shape
+            == torch.stack([torch.empty(shape), torch.empty(shape)], stack_dim).shape
+        )
 
     def test_stack_onehot(self, shape, stack_dim):
         n = 5

torchrl/data/tensor_specs.py

@@ -3683,52 +3683,73 @@ class Choice(TensorSpec):
     """A discrete choice spec for either tensor or non-tensor data.
 
     Args:
-        stack (:class:`~Stacked`, :class:`~StackedComposite`, or :class:`~tensordict.NonTensorStack`):
-            Stack of specs or non-tensor data from which to choose during
-            sampling.
-        device (str, int or torch.device, optional): device of the tensors.
+        choices (list[:class:`~TensorSpec`, :class:`~tensordict.NonTensorData`, :class:`~tensordict.NonTensorStack`]):
+            List of specs or non-tensor data from which to choose during
+            sampling. All elements must have the same type, shape, dtype, and
+            device.
 
     Examples:
         >>> import torch
         >>> _ = torch.manual_seed(0)
-        >>> from torchrl.data import Choice, Categorical
-        >>> spec = Choice(torch.stack([
-        ...     Categorical(n=4, shape=(1,)),
-        ...     Categorical(n=4, shape=(2,))]))
-        >>> spec.shape
-        torch.Size([2, -1])
+        >>> from torchrl.data import Choice, Bounded
+        >>> spec = Choice([
+        ...     Bounded(0, 1, shape=(1,)),
+        ...     Bounded(10, 11, shape=(1,))])
         >>> spec.rand()
-        tensor([3])
+        tensor([0.7682])
         >>> spec.rand()
-        tensor([0, 3])
+        tensor([10.1320])
+        >>> from tensordict import NonTensorData
+        >>> _ = torch.manual_seed(0)
+        >>> spec = Choice([NonTensorData(s) for s in ["a", "b", "c", "d"]])
+        >>> spec.rand().data
+        'a'
+        >>> spec.rand().data
+        'd'
     """
 
     def __init__(
         self,
-        stack: Stacked | StackedComposite | NonTensorStack,
-        device: Optional[DEVICE_TYPING] = None,
+        choices: List[TensorSpec | NonTensorData | NonTensorStack],
     ):
-        assert isinstance(stack, (Stacked, StackedComposite, NonTensorStack))
-        stack = stack.clone()
-        if device is not None:
-            self._stack = stack.to(device)
-        else:
-            self._stack = stack
-            device = stack.device
+        if not isinstance(choices, list):
+            raise TypeError("'choices' must be a list")
 
-        shape = stack.shape
-        dtype = stack.dtype
+        if not isinstance(choices[0], (TensorSpec, NonTensorData, NonTensorStack)):
+            raise TypeError(
+                (
+                    "Each choice must be either a TensorSpec, NonTensorData, or "
+                    f"NonTensorStack, but got {type(choices[0])}"
+                )
+            )
+
+        if not all([isinstance(choice, type(choices[0])) for choice in choices[1:]]):
+            raise TypeError("All choices must be the same type")
+
+        if not all([choice.shape == choices[0].shape for choice in choices[1:]]):
+            raise ValueError("All choices must have the same shape")
+
+        if not all([choice.dtype == choices[0].dtype for choice in choices[1:]]):
+            raise ValueError("All choices must have the same dtype")
+
+        if not all([choice.device == choices[0].device for choice in choices[1:]]):
+            raise ValueError("All choices must have the same device")
+
+        shape = choices[0].shape
+        device = choices[0].device
+        dtype = choices[0].dtype
 
-        domain = None
         super().__init__(
-            shape=shape, space=None, device=device, dtype=dtype, domain=domain
+            shape=shape, space=None, device=device, dtype=dtype, domain=None
         )
 
+        self._choices = [choice.clone() for choice in choices]
+
     def _rand_idx(self):
-        return torch.randint(0, len(self._stack), ()).item()
+        return torch.randint(0, len(self._choices), ()).item()
 
     def _sample(self, idx, spec_sample_fn) -> TensorDictBase:
-        res = self._stack[idx]
+        res = self._choices[idx]
         if isinstance(res, TensorSpec):
             return spec_sample_fn(res)
         else:
@@ -3744,32 +3765,32 @@ def rand(self, shape: torch.Size = None) -> TensorDictBase:
         return self._sample(self._rand_idx(), lambda x: x.rand(shape))
 
     def is_in(self, val: torch.Tensor | TensorDictBase) -> bool:
-        if isinstance(self._stack, (Stacked, StackedComposite)):
-            return any([stack_elem.is_in(val) for stack_elem in self._stack])
+        if isinstance(self._choices[0], TensorSpec):
+            return any([choice.is_in(val) for choice in self._choices])
         else:
-            return any([(stack_elem == val).all() for stack_elem in self._stack])
+            return any([(choice == val).all() for choice in self._choices])
 
     def expand(self, *shape):
-        raise NotImplementedError
+        return self.__class__([choice.expand(*shape) for choice in self._choices])
 
     def unsqueeze(self, dim: int):
-        raise NotImplementedError
+        return self.__class__([choice.unsqueeze(dim) for choice in self._choices])
 
     def clone(self) -> Choice:
-        return self.__class__(self._stack)
+        return self.__class__([choice.clone() for choice in self._choices])
 
     def cardinality(self) -> int:
-        if isinstance(self._stack, NonTensorStack):
-            return len(self._stack)
+        if isinstance(self._choices[0], (NonTensorData, NonTensorStack)):
+            return len(self._choices)
         else:
             return (
-                torch.tensor([stack_elem.cardinality() for stack_elem in self._stack])
+                torch.tensor([choice.cardinality() for choice in self._choices])
                 .sum()
                 .item()
             )
 
     def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> Choice:
-        return self.__class__(self._stack.to(dest))
+        return self.__class__([choice.to(dest) for choice in self._choices])
 
 
 @dataclass(repr=False)