[Refactor] Box device (#881)

Author: vmoens

test/_utils_internal.py

@@ -3,6 +3,7 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+import contextlib
 import os
 import time
 from functools import wraps
@@ -104,3 +105,13 @@ def dtype_fixture():
     torch.set_default_dtype(torch.double)
     yield dtype
     torch.set_default_dtype(dtype)
+
+
+@contextlib.contextmanager
+def set_global_var(module, var_name, value):
+    old_value = getattr(module, var_name)
+    setattr(module, var_name, value)
+    try:
+        yield
+    finally:
+        setattr(module, var_name, old_value)

test/test_specs.py

@@ -7,7 +7,8 @@
 import numpy as np
 import pytest
 import torch
-from _utils_internal import get_available_devices
+import torchrl.data.tensor_specs
+from _utils_internal import get_available_devices, set_global_var
 from scipy.stats import chisquare
 from tensordict.tensordict import TensorDict, TensorDictBase
 from torchrl.data.tensor_specs import (
@@ -57,8 +58,11 @@ def test_discrete(cls):
         ts.encode(torch.tensor([5]))
         ts.encode(torch.tensor(5).numpy())
         ts.encode(9)
-        with pytest.raises(AssertionError):
+        with pytest.raises(AssertionError), set_global_var(
+            torchrl.data.tensor_specs, "_CHECK_SPEC_ENCODE", True
+        ):
             ts.encode(torch.tensor([11]))  # out of bounds
+        assert not torchrl.data.tensor_specs._CHECK_SPEC_ENCODE
         assert ts.is_in(r)
         assert (ts.encode(ts.to_numpy(r)) == r).all()
 
@@ -114,10 +118,15 @@ def test_ndbounded(dtype, shape):
         ts.encode(lb + torch.rand(10) * (ub - lb))
         ts.encode((lb + torch.rand(10) * (ub - lb)).numpy())
         assert (ts.encode(ts.to_numpy(r)) == r).all()
-        with pytest.raises(AssertionError):
+        with pytest.raises(AssertionError), set_global_var(
+            torchrl.data.tensor_specs, "_CHECK_SPEC_ENCODE", True
+        ):
             ts.encode(torch.rand(10) + 3)  # out of bounds
-        with pytest.raises(AssertionError):
+        with pytest.raises(AssertionError), set_global_var(
+            torchrl.data.tensor_specs, "_CHECK_SPEC_ENCODE", True
+        ):
             ts.to_numpy(torch.rand(10) + 3)  # out of bounds
+        assert not torchrl.data.tensor_specs._CHECK_SPEC_ENCODE
 
 
 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.float64, None])

torchrl/data/tensor_specs.py

@@ -35,7 +35,9 @@
 
 INDEX_TYPING = Union[int, torch.Tensor, np.ndarray, slice, List]
 
-_NO_CHECK_SPEC_ENCODE = get_binary_env_var("NO_CHECK_SPEC_ENCODE")
+# By default, we do not check that an obs is in the domain. THis should be done when validating the env beforehand
+_CHECK_SPEC_ENCODE = get_binary_env_var("CHECK_SPEC_ENCODE")
+
 
 _DEFAULT_SHAPE = torch.Size((1,))
 
@@ -108,8 +110,28 @@ def clone(self) -> DiscreteBox:
 class ContinuousBox(Box):
     """A continuous box of values, in between a minimum and a maximum."""
 
-    minimum: torch.Tensor
-    maximum: torch.Tensor
+    _minimum: torch.Tensor
+    _maximum: torch.Tensor
+    device: torch.device = None
+
+    # We store the tensors on CPU to avoid overloading CUDA with tensors that are rarely used.
+    @property
+    def minimum(self):
+        return self._minimum.to(self.device)
+
+    @property
+    def maximum(self):
+        return self._maximum.to(self.device)
+
+    @minimum.setter
+    def minimum(self, value):
+        self.device = value.device
+        self._minimum = value.cpu()
+
+    @maximum.setter
+    def maximum(self, value):
+        self.device = value.device
+        self._maximum = value.cpu()
 
     def __post_init__(self):
         self.minimum = self.minimum.clone()
@@ -257,7 +279,7 @@ def encode(self, val: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
                         f"Shape mismatch: the value has shape {val.shape} which "
                         f"is incompatible with the spec shape {self.shape}."
                     )
-        if not _NO_CHECK_SPEC_ENCODE:
+        if _CHECK_SPEC_ENCODE:
             self.assert_is_in(val)
         return val
 

torchrl/envs/common.py

@@ -210,6 +210,7 @@ class EnvBase(nn.Module, metaclass=abc.ABCMeta):
         - run_type_checks (bool): if True, the observation and reward dtypes
             will be compared against their respective spec and an exception
             will be raised if they don't match.
+            Defaults to False.
 
     Methods:
         step (TensorDictBase -> TensorDictBase): step in the environment
@@ -226,7 +227,7 @@ def __init__(
         device: DEVICE_TYPING = "cpu",
         dtype: Optional[Union[torch.dtype, np.dtype]] = None,
         batch_size: Optional[torch.Size] = None,
-        run_type_checks: bool = True,
+        run_type_checks: bool = False,
     ):
         super().__init__()
         if device is not None:

torchrl/envs/transforms/transforms.py

@@ -2738,7 +2738,9 @@ def transform_observation_spec(
             dtype=torch.int64,
             device=observation_spec.device,
         )
-        observation_spec["step_count"].space.minimum = 0
+        observation_spec["step_count"].space.minimum = (
+            observation_spec["step_count"].space.minimum * 0
+        )
         return observation_spec