Update

[ghstack-poisoned]
pytorch · Aug 6, 2024 · 9f8d0fc · 9f8d0fc
2 parents 4dbd271 + d2af456
commit 9f8d0fc
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
@@ -46,7 +46,7 @@
 from torch.utils._pytree import tree_map
 from torchrl._utils import _replace_last, _rng_decorator, logger as torchrl_logger
 
-from torchrl.data.tensor_specs import Composite, NO_DEFAULT, TensorSpec, Unbounded
+from torchrl.data.tensor_specs import Composite, NO_DEFAULT_RL as NO_DEFAULT, TensorSpec, Unbounded
 from torchrl.data.utils import check_no_exclusive_keys
 
 __all__ = [