TensorSpec.encode domain check parametrization (#228)

Author: vmoens

torchrl/data/tensor_specs.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from textwrap import indent
 from typing import (
@@ -41,6 +42,17 @@
 
 INDEX_TYPING = Union[int, torch.Tensor, np.ndarray, slice, List]
 
+_NO_CHECK_SPEC_ENCODE = os.environ.get("NO_CHECK_SPEC_ENCODE", False)
+if _NO_CHECK_SPEC_ENCODE in ("0", "False", False):
+    _NO_CHECK_SPEC_ENCODE = False
+elif _NO_CHECK_SPEC_ENCODE in ("1", "True", True):
+    _NO_CHECK_SPEC_ENCODE = True
+else:
+    raise NotImplementedError(
+        "NO_CHECK_SPEC_ENCODE should be in 'True', 'False', '0' or '1'. "
+        f"Got {_NO_CHECK_SPEC_ENCODE} instead."
+    )
+
 
 def _default_dtype_and_device(
     dtype: Union[None, torch.dtype],
@@ -214,7 +226,8 @@ def encode(self, val: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
             ):
                 val = val.copy()
             val = torch.as_tensor(val, dtype=self.dtype, device=self.device)
-        self.assert_is_in(val)
+        if not _NO_CHECK_SPEC_ENCODE:
+            self.assert_is_in(val)
         return val
 
     def to_numpy(self, val: torch.Tensor, safe: bool = True) -> np.ndarray: