[BugFix] Fix run_type_checks (#1570)

Author: Vincent Moens

test/test_env.py

@@ -2117,6 +2117,29 @@ def test_num_threads():
         torch.set_num_threads(num_threads)
 
 
+def test_run_type_checks():
+    env = ContinuousActionVecMockEnv()
+    env._run_type_checks = False
+    check_env_specs(env)
+    env._run_type_checks = True
+    check_env_specs(env)
+    env.output_spec.unlock_()
+    # check type check on done
+    env.output_spec["full_done_spec", "done"].dtype = torch.int
+    with pytest.raises(TypeError, match="expected done.dtype to"):
+        check_env_specs(env)
+    env.output_spec["full_done_spec", "done"].dtype = torch.bool
+    # check type check on reward
+    env.output_spec["full_reward_spec", "reward"].dtype = torch.int
+    with pytest.raises(TypeError, match="expected"):
+        check_env_specs(env)
+    env.output_spec["full_reward_spec", "reward"].dtype = torch.float
+    # check type check on obs
+    env.output_spec["full_observation_spec", "observation"].dtype = torch.float16
+    with pytest.raises(TypeError):
+        check_env_specs(env)
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/envs/common.py

@@ -1193,10 +1193,9 @@ def _step_proc_data(self, next_tensordict_out):
                 next_tensordict_out.set(done_key, done)
 
         if self.run_type_checks:
-            # TODO: check these errors
-            for key in self._select_observation_keys(next_tensordict_out):
+            for key, spec in self.observation_spec.items():
                 obs = next_tensordict_out.get(key)
-                self.observation_spec.type_check(obs, key)
+                spec.type_check(obs)
 
             for reward_key in self.reward_keys:
                 if (
@@ -1213,10 +1212,10 @@ def _step_proc_data(self, next_tensordict_out):
             for done_key in self.done_keys:
                 if (
                     next_tensordict_out.get(done_key).dtype
-                    is not self.output_spec["full_done_spec"].get(done_key).dtype
+                    is not self.output_spec["full_done_spec", done_key].dtype
                 ):
                     raise TypeError(
-                        f"expected done.dtype to be torch.bool but got {next_tensordict_out.get(done_key).dtype}"
+                        f"expected done.dtype to be {self.output_spec['full_done_spec', done_key].dtype} but got {next_tensordict_out.get(done_key).dtype}"
                     )
         return next_tensordict_out