diff --git a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
index 729452aff1..23a29f7840 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplPrivateSpaceTest.java
@@ -49,6 +49,7 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
 
     private static final String PRIVATE_PROFILE_NAME = "LauncherPrivateProfile";
     private static final String INSTALLED_APP_NAME = "Aardwolf";
+    private static final int MAX_STATE_TOGGLE_TRIES = 2;
     private static final String TAG = "TaplPrivateSpaceTest";
 
     @Override
@@ -172,7 +173,9 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
         HomeAllApps homeAllApps = mLauncher.getAllApps();
 
         // Disable Private Space
-        togglePrivateSpace(PrivateProfileManager.STATE_DISABLED, homeAllApps);
+        togglePrivateSpaceWithRetry(PrivateProfileManager.STATE_DISABLED, homeAllApps);
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
 
         homeAllApps.freeze();
         try {
@@ -184,7 +187,9 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
         }
 
         // Enable Private Space
-        togglePrivateSpace(PrivateProfileManager.STATE_ENABLED, homeAllApps);
+        togglePrivateSpaceWithRetry(PrivateProfileManager.STATE_ENABLED, homeAllApps);
+        // Scroll to the bottom of All Apps
+        executeOnLauncher(launcher -> launcher.getAppsView().resetAndScrollToPrivateSpaceHeader());
 
         homeAllApps.freeze();
         try {
@@ -217,6 +222,25 @@ public class TaplPrivateSpaceTest extends AbstractQuickStepTest {
         waitForLauncherUIUpdate();
     }
 
+    private void togglePrivateSpaceWithRetry(int state, HomeAllApps homeAllApps) {
+        int togglePsCount = 0;
+        boolean shouldRetry;
+        do {
+            togglePsCount ++;
+            try {
+                togglePrivateSpace(state, homeAllApps);
+                // No need to retry if the toggle was successful.
+                shouldRetry = false;
+            } catch (AssertionError error) {
+                if (togglePsCount < MAX_STATE_TOGGLE_TRIES) {
+                    shouldRetry = true;
+                } else {
+                    throw error;
+                }
+            }
+        } while (shouldRetry);
+    }
+
     private void waitForPrivateSpaceSetup() {
         waitForLauncherCondition("Private Profile not setup",
                 launcher -> launcher.getAppsView().hasPrivateProfile(),