diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index ba6f165..abf49eb 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.views.RecentsView;
 
 import java.io.PrintWriter;
+import java.util.Arrays;
 
 /**
  * A data source which integrates with a Launcher instance
@@ -105,6 +106,10 @@
 
         // 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 9a9e0ba..9c463cb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -923,6 +923,15 @@
         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 43feec7..a1390ae 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -591,9 +591,7 @@
      * 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 66c2eb3..d3f80e3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -196,6 +196,19 @@
         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().