init

pytorch · Aug 6, 2024 · fdb8e12 · fdb8e12
1 parent 76213f7
commit fdb8e12
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import abc
+import enum
 import math
 import warnings
 from collections.abc import Iterable
@@ -81,6 +82,15 @@
 )
 
 
+# Akin to TD's NO_DEFAULT but won't raise a KeyError when found in a TD or used as default
+class _NoDefault(enum.IntEnum):
+    ZERO = 0
+    ONE = 1
+
+
+NO_DEFAULT_RL = _NoDefault.ONE
+
+
 def _default_dtype_and_device(
     dtype: Union[None, torch.dtype],
     device: Union[None, str, int, torch.device],

diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
@@ -48,7 +48,7 @@
 
 from torchrl.data.tensor_specs import (
     CompositeSpec,
-    NO_DEFAULT,
+    NO_DEFAULT_RL as NO_DEFAULT,
     TensorSpec,
     UnboundedContinuousTensorSpec,
 )