[Feature] Add Choice spec

ghstack-source-id: 6776395
Pull Request resolved: #2713

Author: kurtamohler

test/test_specs.py

@@ -12,7 +12,13 @@
 import torch
 import torchrl.data.tensor_specs
 from scipy.stats import chisquare
-from tensordict import LazyStackedTensorDict, TensorDict, TensorDictBase
+from tensordict import (
+    LazyStackedTensorDict,
+    NonTensorData,
+    NonTensorStack,
+    TensorDict,
+    TensorDictBase,
+)
 from tensordict.utils import _unravel_key_to_tuple
 from torchrl._utils import _make_ordinal_device
 
@@ -23,6 +29,7 @@
     Bounded,
     BoundedTensorSpec,
     Categorical,
+    Choice,
     Composite,
     CompositeSpec,
     ContinuousBox,
@@ -702,6 +709,62 @@ def test_change_batch_size(self, shape, is_complete, device, dtype):
         assert ts["nested"].shape == (3,)
 
 
+class TestChoiceSpec:
+    @pytest.mark.parametrize("input_type", ["spec", "nontensor", "nontensorstack"])
+    def test_choice(self, input_type):
+        if input_type == "spec":
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
+            example_in = torch.tensor(11.0)
+            example_out = torch.tensor(9.0)
+        elif input_type == "nontensor":
+            choices = [NonTensorData("a"), NonTensorData("b")]
+            example_in = NonTensorData("b")
+            example_out = NonTensorData("c")
+        elif input_type == "nontensorstack":
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
+            example_in = NonTensorStack("a", "b", "c")
+            example_out = NonTensorStack("a", "c", "b")
+
+        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)])
 @pytest.mark.parametrize("device", get_default_devices())
 def test_create_composite_nested(shape, device):
@@ -1498,6 +1561,27 @@ def test_non_tensor(self):
         )
         assert spec.expand(2, 3, 4).example_data == "example_data"
 
+    @pytest.mark.parametrize("input_type", ["spec", "nontensor", "nontensorstack"])
+    def test_choice(self, input_type):
+        if input_type == "spec":
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
+        elif input_type == "nontensor":
+            choices = [NonTensorData("a"), NonTensorData("b")]
+        elif input_type == "nontensorstack":
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
+
+        spec = Choice(choices)
+        res = spec.expand(
+            [
+                3,
+            ]
+        )
+        assert res.shape == torch.Size(
+            [
+                3,
+            ]
+        )
+
     @pytest.mark.parametrize("shape1", [None, (), (5,)])
     @pytest.mark.parametrize("shape2", [(), (10,)])
     def test_onehot(self, shape1, shape2):
@@ -1701,6 +1785,19 @@ def test_non_tensor(self):
         assert spec.clone() is not spec
         assert spec.clone().example_data == "example_data"
 
+    @pytest.mark.parametrize("input_type", ["spec", "nontensor", "nontensorstack"])
+    def test_choice(self, input_type):
+        if input_type == "spec":
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
+        elif input_type == "nontensor":
+            choices = [NonTensorData("a"), NonTensorData("b")]
+        elif input_type == "nontensorstack":
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
+
+        spec = Choice(choices)
+        assert spec.clone() == spec
+        assert spec.clone() is not spec
+
     @pytest.mark.parametrize("shape1", [None, (), (5,)])
     def test_onehot(
         self,
@@ -1786,6 +1883,31 @@ def test_non_tensor(self):
         with pytest.raises(RuntimeError, match="Cannot enumerate a NonTensorSpec."):
             spec.cardinality()
 
+    @pytest.mark.parametrize(
+        "input_type",
+        ["bounded_spec", "categorical_spec", "nontensor", "nontensorstack"],
+    )
+    def test_choice(self, input_type):
+        if input_type == "bounded_spec":
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
+        elif input_type == "categorical_spec":
+            choices = [Categorical(10), Categorical(20)]
+        elif input_type == "nontensor":
+            choices = [NonTensorData("a"), NonTensorData("b"), NonTensorData("c")]
+        elif input_type == "nontensorstack":
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
+
+        spec = Choice(choices)
+
+        if input_type == "bounded_spec":
+            assert spec.cardinality() == float("inf")
+        elif input_type == "categorical_spec":
+            assert spec.cardinality() == 30
+        elif input_type == "nontensor":
+            assert spec.cardinality() == 3
+        elif input_type == "nontensorstack":
+            assert spec.cardinality() == 2
+
     @pytest.mark.parametrize("shape1", [(5,), (5, 6)])
     def test_onehot(
         self,
@@ -2096,6 +2218,23 @@ def test_non_tensor(self, device):
         assert spec.to(device).device == device
         assert spec.to(device).example_data == "example_data"
 
+    @pytest.mark.parametrize(
+        "input_type",
+        ["bounded_spec", "categorical_spec", "nontensor", "nontensorstack"],
+    )
+    def test_choice(self, input_type, device):
+        if input_type == "bounded_spec":
+            choices = [Bounded(0, 2.5, ()), Bounded(10, 12, ())]
+        elif input_type == "categorical_spec":
+            choices = [Categorical(10), Categorical(20)]
+        elif input_type == "nontensor":
+            choices = [NonTensorData("a"), NonTensorData("b")]
+        elif input_type == "nontensorstack":
+            choices = [NonTensorStack("a", "b", "c"), NonTensorStack("d", "e", "f")]
+
+        spec = Choice(choices)
+        assert spec.to(device).device == device
+
     @pytest.mark.parametrize("shape1", [(5,), (5, 6)])
     def test_onehot(self, shape1, device):
         if shape1 is None:
@@ -2363,6 +2502,33 @@ def test_stack_non_tensor(self, shape, stack_dim):
         assert new_spec.device == torch.device("cpu")
         assert new_spec.example_data == "example_data"
 
+    @pytest.mark.parametrize(
+        "input_type",
+        ["bounded_spec", "categorical_spec", "nontensor"],
+    )
+    def test_stack_choice(self, input_type, shape, stack_dim):
+        if input_type == "bounded_spec":
+            choices = [Bounded(0, 2.5, shape), Bounded(10, 12, shape)]
+        elif input_type == "categorical_spec":
+            choices = [Categorical(10, shape), Categorical(20, shape)]
+        elif input_type == "nontensor":
+            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(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
         shape = (*shape, 5)

torchrl/data/__init__.py

@@ -76,6 +76,7 @@
     BoundedContinuous,
     BoundedTensorSpec,
     Categorical,
+    Choice,
     Composite,
     CompositeSpec,
     DEVICE_TYPING,

torchrl/data/tensor_specs.py

@@ -38,6 +38,7 @@
     is_tensor_collection,
     LazyStackedTensorDict,
     NonTensorData,
+    NonTensorStack,
     TensorDict,
     TensorDictBase,
     unravel_key,
@@ -3732,6 +3733,120 @@ def clone(self) -> Categorical:
         )
 
 
+class Choice(TensorSpec):
+    """A discrete choice spec for either tensor or non-tensor data.
+
+    Args:
+        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, Bounded
+        >>> spec = Choice([
+        ...     Bounded(0, 1, shape=(1,)),
+        ...     Bounded(10, 11, shape=(1,))])
+        >>> spec.rand()
+        tensor([0.7682])
+        >>> spec.rand()
+        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,
+        choices: List[TensorSpec | NonTensorData | NonTensorStack],
+    ):
+        if not isinstance(choices, list):
+            raise TypeError("'choices' must be a list")
+
+        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
+
+        super().__init__(
+            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._choices), ()).item()
+
+    def _sample(self, idx, spec_sample_fn) -> TensorDictBase:
+        res = self._choices[idx]
+        if isinstance(res, TensorSpec):
+            return spec_sample_fn(res)
+        else:
+            return res
+
+    def zero(self, shape: torch.Size = None) -> TensorDictBase:
+        return self._sample(0, lambda x: x.zero(shape))
+
+    def one(self, shape: torch.Size = None) -> TensorDictBase:
+        return self._sample(min(1, len(self - 1)), lambda x: x.one(shape))
+
+    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._choices[0], TensorSpec):
+            return any([choice.is_in(val) for choice in self._choices])
+        else:
+            return any([(choice == val).all() for choice in self._choices])
+
+    def expand(self, *shape):
+        return self.__class__([choice.expand(*shape) for choice in self._choices])
+
+    def unsqueeze(self, dim: int):
+        return self.__class__([choice.unsqueeze(dim) for choice in self._choices])
+
+    def clone(self) -> Choice:
+        return self.__class__([choice.clone() for choice in self._choices])
+
+    def cardinality(self) -> int:
+        if isinstance(self._choices[0], (NonTensorData, NonTensorStack)):
+            return len(self._choices)
+        else:
+            return (
+                torch.tensor([choice.cardinality() for choice in self._choices])
+                .sum()
+                .item()
+            )
+
+    def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> Choice:
+        return self.__class__([choice.to(dest) for choice in self._choices])
+
+
 @dataclass(repr=False)
 class Binary(Categorical):
     """A binary discrete tensor spec.