diff --git a/nasim/envs/environment.py b/nasim/envs/environment.py
index dcd74be4..65d4481d 100755
--- a/nasim/envs/environment.py
+++ b/nasim/envs/environment.py
@@ -424,7 +424,7 @@ def get_action_mask(self):
         mask = np.zeros(self.action_space.n, dtype=np.int64)
         for a_idx in range(self.action_space.n):
             action = self.action_space.get_action(a_idx)
-            if self.network.host_discovered(action.target):
+            if self.current_state.host_discovered(action.target):
                 mask[a_idx] = 1
         return mask
 
diff --git a/nasim/scenarios/loader.py b/nasim/scenarios/loader.py
index 16b5c3ae..05f4c9a9 100755
--- a/nasim/scenarios/loader.py
+++ b/nasim/scenarios/loader.py
@@ -310,7 +310,7 @@ def _validate_single_exploit(self, e_name, e):
             (f"{e_name}. Exploit target OS is invalid. '{e[u.EXPLOIT_OS]}'."
              " Should be None or one of the OS in the os list.")
 
-        assert 0 <= e[u.EXPLOIT_PROB] < 1, \
+        assert 0 <= e[u.EXPLOIT_PROB] <= 1.0, \
             (f"{e_name}. Exploit probability, '{e[u.EXPLOIT_PROB]}' not "
              "a valid probability")