diff --git a/desktop_env/controllers/python.py b/desktop_env/controllers/python.py
index b3dab4a..f80ce2e 100644
--- a/desktop_env/controllers/python.py
+++ b/desktop_env/controllers/python.py
@@ -144,7 +144,7 @@ class PythonController:
             return
 
         action_type = action["action_type"]
-        parameters = action["parameters"] if "parameters" in action else {}
+        parameters = action["parameters"] if "parameters" in action else {param: action[param] for param in action if param != 'action_type'}
         move_mode = random.choice(
             ["pyautogui.easeInQuad", "pyautogui.easeOutQuad", "pyautogui.easeInOutQuad", "pyautogui.easeInBounce",
              "pyautogui.easeInElastic"])
diff --git a/desktop_env/envs/desktop_env.py b/desktop_env/envs/desktop_env.py
index 6f061f9..dce2e77 100644
--- a/desktop_env/envs/desktop_env.py
+++ b/desktop_env/envs/desktop_env.py
@@ -236,7 +236,7 @@ class DesktopEnv(gym.Env):
         info = {}
 
         # handle the special actions
-        if action in ['WAIT', 'FAIL', 'DONE']:
+        if action in ['WAIT', 'FAIL', 'DONE'] or (type(action) == dict and action['action_type'] in ['WAIT', 'FAIL', 'DONE']):
             if action == 'WAIT':
                 time.sleep(pause)
             elif action == 'FAIL':