diff --git a/nasim/envs/action.py b/nasim/envs/action.py
index 0d655fc6..aa7a990b 100755
--- a/nasim/envs/action.py
+++ b/nasim/envs/action.py
@@ -695,9 +695,9 @@ def get_action(self, action_idx):
         Action
             Corresponding Action object
         """
-        assert isinstance(action_idx, int), \
+        assert isinstance(action_idx, (int, np.integer)), \
             ("When using flat action space, action must be an integer"
-             f" or an Action object: {action_idx} is invalid")
+             f" or an Action object. {type(action_idx)} is invalid")
         return self.actions[action_idx]