amend

pytorch · Apr 23, 2024 · 12dfa65 · 12dfa65
1 parent 3a4ae87
commit 12dfa65
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -929,7 +929,20 @@ def __getattr__(self, attr: str) -> Any:
                 attr
             )  # make sure that appropriate exceptions are raised
         except AttributeError as err:
-            if attr.endswith("_spec"):
+            if attr in (
+                "action_spec",
+                "done_spec",
+                "full_action_spec",
+                "full_done_spec",
+                "full_observation_spec",
+                "full_reward_spec",
+                "full_state_spec",
+                "input_spec",
+                "observation_spec",
+                "output_spec",
+                "reward_spec",
+                "state_spec",
+            ):
                 raise AttributeError(
                     f"Could not get {attr} because an internal error was raised. To find what this error "
                     f"is, call env.transform.transform_<placeholder>_spec(env.base_env.spec)."