[Feature] spec.is_empty(recurse)

ghstack-source-id: faa3b1df5133c77462d6dd013d3854d684cc7e94 Pull Request resolved: #2596
pytorch · Nov 24, 2024 · 097d8ad · 097d8ad · github-actions · Nov 24, 2024
1 parent 152bc81
commit 097d8ad
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -4278,8 +4278,22 @@ def shape(self, value: torch.Size):
                     )
         self._shape = _size(value)
 
-    def is_empty(self):
-        """Whether the composite spec contains specs or not."""
+    def is_empty(self, recurse: bool = False):
+        """Whether the composite spec contains specs or not.
+
+        Args:
+            recurse (bool): whether to recursively assess if the spec is empty.
+                If ``True``, will return ``True`` if there are no leaves. If ``False``
+                (default) will return whether there is any spec defined at the root level.
+
+        """
+        if recurse:
+            for spec in self._specs.values():
+                if spec is None:
+                    continue
+                if isinstance(spec, Composite) and spec.is_empty(recurse=True):
+                    continue
+                return False
         return len(self._specs) == 0
 
     @property