Update

[ghstack-poisoned]

Author: vmoens

test/mocking_classes.py

@@ -617,6 +617,8 @@ def _step(
 
 
 class ContinuousActionVecMockEnv(_MockEnv):
+    adapt_dtype: bool = True
+
     @classmethod
     def __new__(
         cls,
@@ -715,7 +717,12 @@ def _step(
             done = done.any(-1)
         done = reward = done.unsqueeze(-1)
         tensordict.set(
-            "reward", reward.to(self.reward_spec.dtype).expand(self.reward_spec.shape)
+            "reward",
+            reward.to(
+                self.reward_spec.dtype
+                if self.adapt_dtype
+                else torch.get_default_dtype()
+            ).expand(self.reward_spec.shape),
         )
         tensordict.set("done", done)
         tensordict.set("terminated", done)

test/test_env.py

@@ -235,6 +235,7 @@
 class TestEnvBase:
     def test_run_type_checks(self):
         env = ContinuousActionVecMockEnv()
+        env.adapt_dtype = False
         env._run_type_checks = False
         check_env_specs(env)
         env._run_type_checks = True