Improve cloud environment handling and test output validation

- Add explicit pyautogui disabling in cloud environments
- Update test_act_train_e2e.py to use more robust output parsing
- Add detailed error messages for policy output normalization checks

imitate_mouse/run_mouse_policy.py

@@ -27,6 +27,10 @@ def run_policy_eval(args, num_steps=100):
     # Skip GUI interaction in cloud environments
     in_cloud = os.environ.get('CI') or not os.environ.get('DISPLAY')
 
+    if in_cloud and pyautogui is not None:
+        pyautogui = None  # Disable GUI functionality
+        print("Running in cloud environment, disabling mouse control")
+
     device = "cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
 
     # Load checkpoint with proper error handling

tests/integration/test_act_train_e2e.py

@@ -44,11 +44,15 @@ class Args:
     with patch('builtins.print') as mock_print:
         run_policy_eval(Args(), num_steps=5)
         # Verify policy is producing normalized outputs
-        outputs = [ast.literal_eval(call[0][0].split(": ")[1]) for call in mock_print.call_args_list if 'Policy output' in call[0][0]]
-        assert len(outputs) > 3, "Insufficient policy predictions"
+        outputs = [
+            ast.literal_eval(call.args[0].split(": ")[1])
+            for call in mock_print.call_args_list
+            if 'Policy output' in call.args[0]
+        ]
+        assert len(outputs) > 3, f"Insufficient policy predictions, got {len(outputs)}"
         for out in outputs:
-            assert 0 <= out[0] <= 1, "X output not normalized"
-            assert 0 <= out[1] <= 1, "Y output not normalized"
+            assert 0 <= out[0] <= 1, f"X output {out[0]} not normalized"
+            assert 0 <= out[1] <= 1, f"Y output {out[1]} not normalized"
 
 
 @pytest.mark.integration