bf

pytorch · vmoens · Nov 19, 2022 · Nov 2, 2022 · Nov 7, 2022 · Nov 7, 2022
commit 6dc0a943eaf018b3f07929925562d43040ac3419
diff --git a/test/_utils_internal.py b/test/_utils_internal.py
@@ -50,9 +50,10 @@ def _test_fake_tensordict(env: EnvBase):
     fake_tensordict = fake_tensordict.unsqueeze(real_tensordict.batch_dims - 1)
     fake_tensordict = fake_tensordict.expand(*real_tensordict.shape)
     fake_tensordict = fake_tensordict.to_tensordict()
-    fake_tensordict.zero_()
-    real_tensordict.zero_()
-    assert (fake_tensordict == real_tensordict).all()
+    assert (
+        fake_tensordict.apply(lambda x: torch.zeros_like(x))
+        == real_tensordict.apply(lambda x: torch.zeros_like(x))
+    ).all()
     for key in keys2:
         assert fake_tensordict[key].shape == real_tensordict[key].shape
 
@@ -69,10 +70,10 @@ def _check_dtype(key, value, obs_spec, input_spec):
             _check_dtype(_key, _value, obs_spec, input_spec)
         return
     elif key in input_spec.keys(yield_nesting_keys=True):
-        assert input_spec[key].is_in(value)
+        assert input_spec[key].is_in(value), (input_spec[key], value)
         return
     elif key in obs_spec.keys(yield_nesting_keys=True):
-        assert obs_spec[key].is_in(value)
+        assert obs_spec[key].is_in(value), (input_spec[key], value)
         return
     else:
         raise KeyError(key)

diff --git a/torchrl/envs/common.py b/torchrl/envs/common.py
@@ -646,11 +646,11 @@ def to(self, device: DEVICE_TYPING) -> EnvBase:
     def fake_tensordict(self) -> TensorDictBase:
         """Returns a fake tensordict with key-value pairs that match in shape, device and dtype what can be expected during an environment rollout."""
         input_spec = self.input_spec
-        fake_input = input_spec.zero(self.batch_size)
+        fake_input = input_spec.rand(self.batch_size)
         observation_spec = self.observation_spec
-        fake_obs = observation_spec.zero(self.batch_size)
+        fake_obs = observation_spec.rand(self.batch_size)
         reward_spec = self.reward_spec
-        fake_reward = reward_spec.zero(self.batch_size)
+        fake_reward = reward_spec.rand(self.batch_size)
         fake_td = TensorDict(
             {
                 **fake_obs,