[Feature] Dynamic specs for make_composite_from_td

ghstack-source-id: 79e31e7
Pull Request resolved: #2829

Author: vmoens

test/test_env.py

@@ -16,7 +16,7 @@
 from collections import defaultdict
 from functools import partial
 from sys import platform
-from typing import Optional
+from typing import Any, Optional
 
 import numpy as np
 import pytest
@@ -33,7 +33,7 @@
     TensorDictBase,
 )
 from tensordict.nn import TensorDictModuleBase
-from tensordict.tensorclass import NonTensorStack
+from tensordict.tensorclass import NonTensorStack, TensorClass
 from tensordict.utils import _unravel_key_to_tuple
 from torch import nn
 
@@ -340,7 +340,8 @@ def forward(self, values):
         )
         env.rollout(10, policy)
 
-    def test_make_spec_from_td(self):
+    @pytest.mark.parametrize("dynamic_shape", [True, False])
+    def test_make_spec_from_td(self, dynamic_shape):
         data = TensorDict(
             {
                 "obs": torch.randn(3),
@@ -353,10 +354,44 @@ def test_make_spec_from_td(self):
             },
             [],
         )
-        spec = make_composite_from_td(data)
+        spec = make_composite_from_td(data, dynamic_shape=dynamic_shape)
         assert (spec.zero() == data.zero_()).all()
         for key, val in data.items(True, True):
             assert val.dtype is spec[key].dtype
+        if dynamic_shape:
+            assert all(s.shape[-1] == -1 for s in spec.values(True, True))
+
+    def test_make_spec_from_tc(self):
+        class Scratch(TensorClass):
+            obs: torch.Tensor
+            string: str
+            some_object: Any
+
+        class Whatever:
+            ...
+
+        td = TensorDict(
+            a=Scratch(
+                obs=torch.ones(5, 3),
+                string="another string!",
+                some_object=Whatever(),
+                batch_size=(5,),
+            ),
+            b="a string!",
+            batch_size=(5,),
+        )
+        spec = make_composite_from_td(td)
+        assert isinstance(spec, Composite)
+        assert isinstance(spec["a"], Composite)
+        assert isinstance(spec["b"], NonTensor)
+        assert spec["b"].example_data == "a string!", spec["b"].example_data
+        assert spec["a", "string"].example_data == "another string!"
+        one = spec.one()
+        assert isinstance(one["a"], Scratch)
+        assert isinstance(one["b"], str)
+        assert isinstance(one["a"].string, str)
+        assert isinstance(one["a"].some_object, Whatever)
+        assert (one == td).all()
 
     def test_env_that_does_nothing(self):
         env = EnvThatDoesNothing()

torchrl/envs/utils.py

@@ -888,13 +888,19 @@ def _sort_keys(element):
     return element
 
 
-def make_composite_from_td(data, unsqueeze_null_shapes: bool = True):
+def make_composite_from_td(
+    data, *, unsqueeze_null_shapes: bool = True, dynamic_shape: bool = False
+):
     """Creates a Composite instance from a tensordict, assuming all values are unbounded.
 
     Args:
         data (tensordict.TensorDict): a tensordict to be mapped onto a Composite.
+
+    Keyword Args:
         unsqueeze_null_shapes (bool, optional): if ``True``, every empty shape will be
             unsqueezed to (1,). Defaults to ``True``.
+        dynamic_shape (bool, optional): if ``True``, all tensors will be assumed to have a dynamic shape
+            along the last dimension. Defaults to ``False``.
 
     Examples:
         >>> from tensordict import TensorDict
@@ -919,22 +925,33 @@ def make_composite_from_td(data, unsqueeze_null_shapes: bool = True):
     """
     # custom function to convert a tensordict in a similar spec structure
     # of unbounded values.
+    def make_shape(shape):
+        if shape or not unsqueeze_null_shapes:
+            if dynamic_shape:
+                return shape[:-1] + (-1,)
+            else:
+                return shape
+        return torch.Size([1])
+
     composite = Composite(
         {
-            key: make_composite_from_td(tensor)
-            if isinstance(tensor, TensorDictBase)
-            else NonTensor(shape=data.shape, device=tensor.device)
+            key: make_composite_from_td(
+                tensor,
+                unsqueeze_null_shapes=unsqueeze_null_shapes,
+                dynamic_shape=dynamic_shape,
+            )
+            if is_tensor_collection(tensor) and not is_non_tensor(tensor)
+            else NonTensor(
+                shape=tensor.shape, example_data=tensor.data, device=tensor.device
+            )
             if is_non_tensor(tensor)
             else Unbounded(
-                dtype=tensor.dtype,
-                device=tensor.device,
-                shape=tensor.shape
-                if tensor.shape or not unsqueeze_null_shapes
-                else [1],
+                dtype=tensor.dtype, device=tensor.device, shape=make_shape(tensor.shape)
             )
             for key, tensor in data.items()
         },
         shape=data.shape,
+        data_cls=type(data),
     )
     return composite