diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index ba6f1651ed..abf49eb195 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -50,6 +50,7 @@ import com.android.quickstep.util.GroupTask;
 import com.android.quickstep.views.RecentsView;
 
 import java.io.PrintWriter;
+import java.util.Arrays;
 
 /**
  * A data source which integrates with a Launcher instance
@@ -105,6 +106,10 @@ public class LauncherTaskbarUIController extends TaskbarUIController {
 
         // Restore the in-app display progress from before Taskbar was recreated.
         float[] prevProgresses = mControllers.getSharedState().inAppDisplayProgressMultiPropValues;
+        // Make a copy of the previous progress to set since updating the multiprop will update
+        // the property which also calls onInAppDisplayProgressChanged() which writes the current
+        // values into the shared state
+        prevProgresses = Arrays.copyOf(prevProgresses, prevProgresses.length);
         for (int i = 0; i < prevProgresses.length; i++) {
             mTaskbarInAppDisplayProgressMultiProp.get(i).setValue(prevProgresses[i]);
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 9a9e0ba70c..9c463cb68c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -923,6 +923,15 @@ public class NavbarButtonsViewController implements TaskbarControllers.LoggableT
         insetsInfo.setTouchableInsets(TOUCHABLE_INSETS_REGION);
     }
 
+    /**
+     * Called whenever a new ui controller is set, and should update anything that depends on the
+     * ui controller.
+     */
+    public void onUiControllerChanged() {
+        updateNavButtonInAppDisplayProgressForSysui();
+        updateNavButtonTranslationY();
+    }
+
     @Override
     public void dumpLogs(String prefix, PrintWriter pw) {
         pw.println(prefix + "NavbarButtonsViewController:");
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 43feec716d..a1390aeabb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -591,9 +591,7 @@ public class TaskbarActivityContext extends BaseTaskbarContext {
      * Sets a new data-source for this taskbar instance
      */
     public void setUIController(@NonNull TaskbarUIController uiController) {
-        mControllers.uiController.onDestroy();
-        mControllers.uiController = uiController;
-        mControllers.uiController.init(mControllers);
+        mControllers.setUiController(uiController);
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 66c2eb3314..d3f80e31cd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -196,6 +196,19 @@ public class TaskbarControllers {
         mPostInitCallbacks.clear();
     }
 
+    /**
+     * Sets the ui controller.
+     */
+    public void setUiController(@NonNull TaskbarUIController newUiController) {
+        uiController.onDestroy();
+        uiController = newUiController;
+        uiController.init(this);
+        uiController.updateStateForSysuiFlags(mSharedState.sysuiStateFlags);
+
+        // Notify that the ui controller has changed
+        navbarButtonsViewController.onUiControllerChanged();
+    }
+
     @Nullable
     public TaskbarSharedState getSharedState() {
         // This should only be null if called before init() and after destroy().