[BugFix] Fix flaky gym penv test (#1853)

Author: Vincent Moens

test/_utils_internal.py

@@ -330,9 +330,9 @@ def rollout_consistency_assertion(
 ):
     """Tests that observations in "next" match observations in the next root tensordict when done is False, and don't match otherwise."""
 
-    done = rollout[:, :-1]["next", done_key].squeeze(-1)
+    done = rollout[..., :-1]["next", done_key].squeeze(-1)
     # data resulting from step, when it's not done
-    r_not_done = rollout[:, :-1]["next"][~done]
+    r_not_done = rollout[..., :-1]["next"][~done]
     # data resulting from step, when it's not done, after step_mdp
     r_not_done_tp1 = rollout[:, 1:][~done]
     torch.testing.assert_close(
@@ -343,17 +343,15 @@ def rollout_consistency_assertion(
 
     if done_strict and not done.any():
         raise RuntimeError("No done detected, test could not complete.")
-
-    # data resulting from step, when it's done
-    r_done = rollout[:, :-1]["next"][done]
-    # data resulting from step, when it's done, after step_mdp and reset
-    r_done_tp1 = rollout[:, 1:][done]
-    assert (
-        (r_done[observation_key] - r_done_tp1[observation_key]).norm(dim=-1) > 1e-1
-    ).all(), (
-        f"Entries in next tensordict do not match entries in root "
-        f"tensordict after reset : {(r_done[observation_key] - r_done_tp1[observation_key]).norm(dim=-1) < 1e-1}"
-    )
+    if done.any():
+        # data resulting from step, when it's done
+        r_done = rollout[..., :-1]["next"][done]
+        # data resulting from step, when it's done, after step_mdp and reset
+        r_done_tp1 = rollout[..., 1:][done]
+        # check that at least one obs after reset does not match the version before reset
+        assert not torch.isclose(
+            r_done[observation_key], r_done_tp1[observation_key]
+        ).all()
 
 
 def rand_reset(env):

torchrl/data/datasets/minari_data.py

@@ -412,8 +412,8 @@ def _proc_spec(spec):
             )
         return BoundedTensorSpec(
             shape=spec["shape"],
-            low=torch.tensor(spec["low"]),
-            high=torch.tensor(spec["high"]),
+            low=torch.as_tensor(spec["low"]),
+            high=torch.as_tensor(spec["high"]),
             dtype=_DTYPE_DIR[spec["dtype"]],
         )
     elif spec["type"] == "Discrete":

torchrl/data/datasets/openx.py

@@ -684,7 +684,7 @@ def _slice_data(data: TensorDict, slice_len, pad_value):
         truncated,
         dim=data.ndim - 1,
         value=True,
-        index=torch.tensor(-1, device=truncated.device),
+        index=torch.as_tensor(-1, device=truncated.device),
     )
     done = data.get(("next", "done"))
     data.set(("next", "truncated"), truncated)

torchrl/data/replay_buffers/replay_buffers.py

@@ -867,7 +867,7 @@ def add(self, data: TensorDictBase) -> int:
                 device=data.device,
             )
             if data.batch_size:
-                data_add["_rb_batch_size"] = torch.tensor(data.batch_size)
+                data_add["_rb_batch_size"] = torch.as_tensor(data.batch_size)
 
         else:
             data_add = data
@@ -1441,7 +1441,7 @@ def __getitem__(
         if isinstance(index, slice) and index == slice(None):
             return self
         if isinstance(index, (list, range, np.ndarray)):
-            index = torch.tensor(index)
+            index = torch.as_tensor(index)
         if isinstance(index, torch.Tensor):
             if index.ndim > 1:
                 raise RuntimeError(

torchrl/data/replay_buffers/samplers.py

@@ -461,10 +461,10 @@ def dumps(self, path):
                 filename=path / "mintree.memmap",
             )
         mm_st.copy_(
-            torch.tensor([self._sum_tree[i] for i in range(self._max_capacity)])
+            torch.as_tensor([self._sum_tree[i] for i in range(self._max_capacity)])
         )
         mm_mt.copy_(
-            torch.tensor([self._min_tree[i] for i in range(self._max_capacity)])
+            torch.as_tensor([self._min_tree[i] for i in range(self._max_capacity)])
         )
         with open(path / "sampler_metadata.json", "w") as file:
             json.dump(

torchrl/data/replay_buffers/storages.py

@@ -1005,7 +1005,7 @@ def __getitem__(self, index):
         if isinstance(index, slice) and index == slice(None):
             return self
         if isinstance(index, (list, range, np.ndarray)):
-            index = torch.tensor(index)
+            index = torch.as_tensor(index)
         if isinstance(index, torch.Tensor):
             if index.ndim > 1:
                 raise RuntimeError(

torchrl/data/replay_buffers/utils.py

@@ -28,11 +28,11 @@ def _to_torch(
     data: Tensor, device, pin_memory: bool = False, non_blocking: bool = False
 ) -> torch.Tensor:
     if isinstance(data, np.generic):
-        return torch.tensor(data, device=device)
+        return torch.as_tensor(data, device=device)
     elif isinstance(data, np.ndarray):
         data = torch.from_numpy(data)
     elif not isinstance(data, Tensor):
-        data = torch.tensor(data, device=device)
+        data = torch.as_tensor(data, device=device)
 
     if pin_memory:
         data = data.pin_memory()

torchrl/data/replay_buffers/writers.py

@@ -357,7 +357,7 @@ def __getstate__(self):
     def dumps(self, path):
         path = Path(path).absolute()
         path.mkdir(exist_ok=True)
-        t = torch.tensor(self._current_top_values)
+        t = torch.as_tensor(self._current_top_values)
         try:
             MemoryMappedTensor.from_filename(
                 filename=path / "current_top_values.memmap",
@@ -453,7 +453,7 @@ def __getitem__(self, index):
         if isinstance(index, slice) and index == slice(None):
             return self
         if isinstance(index, (list, range, np.ndarray)):
-            index = torch.tensor(index)
+            index = torch.as_tensor(index)
         if isinstance(index, torch.Tensor):
             if index.ndim > 1:
                 raise RuntimeError(

torchrl/data/rlhf/utils.py

@@ -100,7 +100,7 @@ def update(self, kl_values: Sequence[float]):
             )
         n_steps = len(kl_values)
         # renormalize kls
-        kl_value = -torch.tensor(kl_values).mean() / self.coef
+        kl_value = -torch.as_tensor(kl_values).mean() / self.coef
         proportional_error = np.clip(kl_value / self.target - 1, -0.2, 0.2)  # ϵₜ
         mult = 1 + proportional_error * n_steps / self.horizon
         self.coef *= mult  # βₜ₊₁
@@ -314,10 +314,10 @@ def _get_done_status(self, generated, batch):
         # of generated tokens
         done_idx = torch.minimum(
             (generated != self.EOS_TOKEN_ID).sum(dim=-1) - batch.prompt_rindex,
-            torch.tensor(self.max_new_tokens) - 1,
+            torch.as_tensor(self.max_new_tokens) - 1,
         )
         truncated_idx = (
-            torch.tensor(self.max_new_tokens, device=generated.device).expand_as(
+            torch.as_tensor(self.max_new_tokens, device=generated.device).expand_as(
                 done_idx
             )
             - 1

torchrl/data/tensor_specs.py

@@ -1374,9 +1374,9 @@ def encode(
     ) -> torch.Tensor:
         if not isinstance(val, torch.Tensor):
             if ignore_device:
-                val = torch.tensor(val)
+                val = torch.as_tensor(val)
             else:
-                val = torch.tensor(val, device=self.device)
+                val = torch.as_tensor(val, device=self.device)
 
         if space is None:
             space = self.space
@@ -1555,9 +1555,9 @@ def __init__(
             dtype = torch.get_default_dtype()
 
         if not isinstance(low, torch.Tensor):
-            low = torch.tensor(low, dtype=dtype, device=device)
+            low = torch.as_tensor(low, dtype=dtype, device=device)
         if not isinstance(high, torch.Tensor):
-            high = torch.tensor(high, dtype=dtype, device=device)
+            high = torch.as_tensor(high, dtype=dtype, device=device)
         if high.device != device:
             high = high.to(device)
         if low.device != device:
@@ -1857,8 +1857,8 @@ def __init__(
         dtype, device = _default_dtype_and_device(dtype, device)
         box = (
             ContinuousBox(
-                torch.tensor(-np.inf, device=device).expand(shape),
-                torch.tensor(np.inf, device=device).expand(shape),
+                torch.as_tensor(-np.inf, device=device).expand(shape),
+                torch.as_tensor(np.inf, device=device).expand(shape),
             )
             if shape == _DEFAULT_SHAPE
             else None