Update the jumpy error to specify the pip install is jax-jumpy (#255)

Farama-Foundation · Jan 11, 2023 · febe526 · febe526
1 parent c52ef43
commit febe526
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/gymnasium/experimental/wrappers/lambda_action.py b/gymnasium/experimental/wrappers/lambda_action.py
@@ -8,7 +8,11 @@
 
 from typing import Callable
 
-import jumpy as jp
+
+try:
+    import jumpy as jp
+except ImportError as e:
+    raise ImportError("Jumpy is not installed, run `pip install jax-jumpy`") from e
 import numpy as np
 
 import gymnasium as gym

diff --git a/gymnasium/experimental/wrappers/lambda_observations.py b/gymnasium/experimental/wrappers/lambda_observations.py
@@ -16,7 +16,11 @@
 from typing import Any, Callable, Sequence
 from typing_extensions import Final
 
-import jumpy as jp
+
+try:
+    import jumpy as jp
+except ImportError as e:
+    raise ImportError("Jumpy is not installed, run `pip install jax-jumpy`") from e
 import numpy as np
 
 import gymnasium as gym

diff --git a/gymnasium/experimental/wrappers/stateful_observation.py b/gymnasium/experimental/wrappers/stateful_observation.py
@@ -11,7 +11,11 @@
 from typing import Any, SupportsFloat
 from typing_extensions import Final
 
-import jumpy as jp
+
+try:
+    import jumpy as jp
+except ImportError as e:
+    raise ImportError("Jumpy is not installed, run `pip install jax-jumpy`") from e
 import numpy as np
 
 import gymnasium as gym