init

pytorch · vmoens · Jun 19, 2024 · Feb 21, 2024 · Feb 21, 2024 · Jun 11, 2024
commit 194831df9773faebe18ce921bdacdd80bea3cf05
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
@@ -16,6 +16,7 @@
     CompositeSpec,
     DiscreteTensorSpec,
     MultiOneHotDiscreteTensorSpec,
+    NonTensorSpec,
     OneHotDiscreteTensorSpec,
     TensorSpec,
     UnboundedContinuousTensorSpec,
@@ -1828,3 +1829,33 @@ def _step(
 
     def _set_seed(self, seed: Optional[int]):
         torch.manual_seed(seed)
+
+
+class EnvWithMetadata(EnvBase):
+    def __init__(self):
+        super().__init__()
+        self.observation_spec = CompositeSpec(
+            tensor=UnboundedContinuousTensorSpec(3),
+            non_tensor=NonTensorSpec(shape=()),
+        )
+        self.reward_spec = UnboundedContinuousTensorSpec(1)
+        self.action_spec = UnboundedContinuousTensorSpec(1)
+
+    def _reset(self, tensordict):
+        data = self.observation_spec.zero()
+        data.set_non_tensor("non_tensor", 0)
+        data.update(self.full_done_spec.zero())
+        return data
+
+    def _step(
+        self,
+        tensordict: TensorDictBase,
+    ) -> TensorDictBase:
+        data = self.observation_spec.zero()
+        data.set_non_tensor("non_tensor", tensordict["non_tensor"] + 1)
+        data.update(self.full_done_spec.zero())
+        data.update(self.full_reward_spec.zero())
+        return data
+
+    def _set_seed(self, seed: Optional[int]):
+        return seed
diff --git a/test/test_env.py b/test/test_env.py
@@ -40,6 +40,7 @@
     DiscreteActionConvMockEnvNumpy,
     DiscreteActionVecMockEnv,
     DummyModelBasedEnvBase,
+    EnvWithMetadata,
     HeterogeneousCountingEnv,
     HeterogeneousCountingEnvPolicy,
     MockBatchedLockedEnv,
@@ -2198,6 +2199,7 @@ def test_parallel(
 @pytest.mark.parametrize(
     "envclass",
     [
+        EnvWithMetadata,
         ContinuousActionConvMockEnv,
         ContinuousActionConvMockEnvNumpy,
         ContinuousActionVecMockEnv,
@@ -2222,6 +2224,7 @@ def test_mocking_envs(envclass):
     env.set_seed(100)
     reset = env.reset()
     _ = env.rand_step(reset)
+    r = env.rollout(3)
     check_env_specs(env, seed=100, return_contiguous=False)
 
 

diff --git a/torchrl/data/__init__.py b/torchrl/data/__init__.py
@@ -51,6 +51,7 @@
     LazyStackedTensorSpec,
     MultiDiscreteTensorSpec,
     MultiOneHotDiscreteTensorSpec,
+    NonTensorSpec,
     OneHotDiscreteTensorSpec,
     TensorSpec,
     UnboundedContinuousTensorSpec,

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -31,7 +31,13 @@
 
 import numpy as np
 import torch
-from tensordict import LazyStackedTensorDict, TensorDict, TensorDictBase, unravel_key
+from tensordict import (
+    LazyStackedTensorDict,
+    NonTensorData,
+    TensorDict,
+    TensorDictBase,
+    unravel_key,
+)
 from tensordict.utils import _getitem_batch_size, NestedKey
 
 from torchrl._utils import get_binary_env_var
@@ -642,7 +648,7 @@ def index(self, index: INDEX_TYPING, tensor_to_index: torch.Tensor) -> torch.Ten
             indexed tensor
 
         """
-        raise NotImplementedError
+        ...
 
     @abc.abstractmethod
     def expand(self, *shape):
@@ -655,7 +661,7 @@ def expand(self, *shape):
                 from it if the current dimension is a singleton.
 
         """
-        raise NotImplementedError
+        ...
 
     def squeeze(self, dim: int | None = None):
         """Returns a new Spec with all the dimensions of size ``1`` removed.
@@ -689,7 +695,7 @@ def is_in(self, val: torch.Tensor) -> bool:
             boolean indicating if values belongs to the TensorSpec box
 
         """
-        raise NotImplementedError
+        ...
 
     def project(self, val: torch.Tensor) -> torch.Tensor:
         """If the input tensor is not in the TensorSpec box, it maps it back to it given some heuristic.
@@ -3952,6 +3958,39 @@ def locked(self):
         return self._locked
 
 
+class NonTensorSpec(TensorSpec):
+    """Tensor spec for non-tensor data.
+
+    This spec has no space or dtype, but the device and shape can be specified
+    (as it is the case for :class:`~tensordict.NonTensorData`).
+
+    """
+
+    shape: torch.Size
+    space: Union[None, Box] = None
+    device: torch.device | None = None
+    dtype: torch.dtype = None
+    domain: str = "continuous"
+
+    def __init__(self, shape=(), device=None):
+        return super().__init__(shape=shape, space=None, device=device, dtype=None)
+
+    def rand(self, shape=None) -> torch.Tensor:
+        return NonTensorData(None, batch_size=self.shape)
+
+    def zero(self, shape=None) -> torch.Tensor:
+        return NonTensorData(None, batch_size=self.shape)
+
+    def expand(self, *shape):
+        return NonTensorSpec(shape=shape)
+
+    def is_in(self, val: torch.Tensor) -> bool:
+        return True
+
+    def index(self, index: INDEX_TYPING, tensor_to_index: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+
 class LazyStackedCompositeSpec(_LazyStackedMixin[CompositeSpec], CompositeSpec):
     """A lazy representation of a stack of composite specs.
 

diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
@@ -5,18 +5,20 @@
 from __future__ import annotations
 
 import contextlib
-
 import importlib.util
 import os
 import re
 from enum import Enum
 from typing import Dict, List, Union
 
+import tensordict.base
+
 import torch
 
 from tensordict import (
     is_tensor_collection,
     LazyStackedTensorDict,
+    NonTensorData,
     TensorDictBase,
     unravel_key,
 )
@@ -290,7 +292,7 @@ def _set(source, dest, key, total_key, excluded):
     if unravel_key(total_key) not in excluded:
         try:
             val = source.get(key)
-            if is_tensor_collection(val):
+            if is_tensor_collection(val) and not isinstance(val, NonTensorData):
                 new_val = dest.get(key, None)
                 if new_val is None:
                     new_val = val.empty()
@@ -471,21 +473,31 @@ def check_env_specs(
             [fake_tensordict.clone() for _ in range(3)], -1
         )
     # eliminate empty containers
-    fake_tensordict_select = fake_tensordict.select(*fake_tensordict.keys(True, True))
-    real_tensordict_select = real_tensordict.select(*real_tensordict.keys(True, True))
+    fake_tensordict_select = fake_tensordict.select(
+        *fake_tensordict.keys(True, True, is_leaf=tensordict.base._default_is_leaf)
+    )
+    real_tensordict_select = real_tensordict.select(
+        *real_tensordict.keys(True, True, is_leaf=tensordict.base._default_is_leaf)
+    )
     # check keys
-    fake_tensordict_keys = set(fake_tensordict.keys(True, True))
-    real_tensordict_keys = set(real_tensordict.keys(True, True))
+    fake_tensordict_keys = set(
+        fake_tensordict.keys(True, True, is_leaf=tensordict.base._is_leaf_nontensor)
+    )
+    real_tensordict_keys = set(
+        real_tensordict.keys(True, True, is_leaf=tensordict.base._is_leaf_nontensor)
+    )
     if fake_tensordict_keys != real_tensordict_keys:
         raise AssertionError(
             f"""The keys of the specs and data do not match:
     - List of keys present in real but not in fake: {real_tensordict_keys-fake_tensordict_keys},
     - List of keys present in fake but not in real: {fake_tensordict_keys-real_tensordict_keys}.
 """
         )
+    print(torch.zeros_like(fake_tensordict_select))
+    print(torch.zeros_like(real_tensordict_select))
     if (
-        fake_tensordict_select.apply(lambda x: torch.zeros_like(x))
-        != real_tensordict_select.apply(lambda x: torch.zeros_like(x))
+        torch.zeros_like(fake_tensordict_select)
+        != torch.zeros_like(real_tensordict_select)
     ).any():
         raise AssertionError(
             "zeroing the two tensordicts did not make them identical. "