diff --git a/quickstep/src/com/android/launcher3/LauncherAppTransitionManagerImpl.java b/quickstep/src/com/android/launcher3/LauncherAppTransitionManagerImpl.java
index b895704..72d6260 100644
--- a/quickstep/src/com/android/launcher3/LauncherAppTransitionManagerImpl.java
+++ b/quickstep/src/com/android/launcher3/LauncherAppTransitionManagerImpl.java
@@ -672,6 +672,7 @@
 
                 if (mLauncher.hasSomeInvisibleFlag(PENDING_INVISIBLE_BY_WALLPAPER_ANIMATION)) {
                     mLauncher.addForceInvisibleFlag(INVISIBLE_BY_PENDING_FLAGS);
+                    mLauncher.getStateManager().moveToRestState();
                 }
 
                 AnimatorSet anim = null;
diff --git a/quickstep/src/com/android/launcher3/uioverrides/BackButtonAlphaHandler.java b/quickstep/src/com/android/launcher3/uioverrides/BackButtonAlphaHandler.java
index 2e6dcc0..722f51b 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/BackButtonAlphaHandler.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/BackButtonAlphaHandler.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.uioverrides;
 
+import android.animation.Animator;
+import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
 
 import com.android.launcher3.Launcher;
@@ -56,6 +58,13 @@
                 final float alpha = (float) valueAnimator.getAnimatedValue();
                 mOverviewInteractionState.setBackButtonAlpha(alpha, false);
             });
+            anim.addListener(new AnimatorListenerAdapter() {
+                @Override
+                public void onAnimationEnd(Animator animation) {
+                    // Reapply the final alpha in case some state (e.g. window focus) changed.
+                    UiFactory.onLauncherStateOrFocusChanged(mLauncher);
+                }
+            });
             builder.play(anim);
         }
     }
diff --git a/quickstep/src/com/android/quickstep/ActivityControlHelper.java b/quickstep/src/com/android/quickstep/ActivityControlHelper.java
index 0205c1f..f1aff05 100644
--- a/quickstep/src/com/android/quickstep/ActivityControlHelper.java
+++ b/quickstep/src/com/android/quickstep/ActivityControlHelper.java
@@ -16,6 +16,7 @@
 package com.android.quickstep;
 
 import static android.view.View.TRANSLATION_Y;
+
 import static com.android.launcher3.LauncherAnimUtils.OVERVIEW_TRANSITION_MS;
 import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
 import static com.android.launcher3.LauncherState.FAST_OVERVIEW;
@@ -52,6 +53,7 @@
 import com.android.launcher3.R;
 import com.android.launcher3.allapps.AllAppsTransitionController;
 import com.android.launcher3.allapps.DiscoveryBounce;
+import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.uioverrides.FastOverviewState;
@@ -265,6 +267,14 @@
                         controller, ALL_APPS_PROGRESS, startProgress, endProgress);
                 shiftAnim.setInterpolator(LINEAR);
                 anim.play(shiftAnim);
+
+                // Since we are changing the start position of the UI, reapply the state, at the end
+                anim.addListener(new AnimationSuccessListener() {
+                    @Override
+                    public void onAnimationSuccess(Animator animator) {
+                        activity.getStateManager().reapplyState();
+                    }
+                });
             }
 
             if (interactionType == INTERACTION_NORMAL) {
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationWrapper.java b/quickstep/src/com/android/quickstep/RecentsAnimationWrapper.java
index 34d42ac..b0313fc 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationWrapper.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationWrapper.java
@@ -49,6 +49,9 @@
         this.mController = controller;
         this.targetSet = targetSet;
 
+        if (controller == null) {
+            return;
+        }
         if (mInputConsumerEnabled) {
             enableInputConsumer();
         }
diff --git a/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java b/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java
index 66bc501..410daa3 100644
--- a/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -1021,8 +1021,6 @@
     private void setTargetAlphaProvider(
             BiFunction<RemoteAnimationTargetCompat, Float, Float> provider) {
         mClipAnimationHelper.setTaskAlphaCallback(provider);
-        // TODO: For some reason, when calling updateFinalShift multiple times on the same frame,
-        // only the first callback is executed.
-        Utilities.postAsyncCallback(mMainThreadHandler, this::updateFinalShift);
+        updateFinalShift();
     }
 }