diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java
index 8cc8965..edcd4c8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarKeyguardController.java
@@ -5,6 +5,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BACK_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BOUNCER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DEVICE_DOZING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DEVICE_DREAMING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_OVERVIEW_DISABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_ON;
@@ -34,7 +35,8 @@
     // locked.
     public static final int MASK_ANY_SYSUI_LOCKED = SYSUI_STATE_BOUNCER_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING
-            | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED;
+            | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED
+            | SYSUI_STATE_DEVICE_DREAMING;
 
     private final TaskbarActivityContext mContext;
     private int mKeyguardSysuiFlags;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index d997f56..cf8148e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -457,8 +457,13 @@
             animatorSet.play(mTaskbarCornerRoundness.animateToValue(cornerRoundness));
         }
 
-        if (hasAnyFlag(changedFlags, FLAG_DEVICE_LOCKED)) {
-            // When transitioning between locked/unlocked, there is no stashing animation.
+        // Keep isUnlockTransition in sync with its counterpart in
+        // TaskbarStashController#createAnimToIsStashed.
+        boolean isUnlockTransition =
+                hasAnyFlag(changedFlags, FLAG_DEVICE_LOCKED) && !hasAnyFlag(FLAG_DEVICE_LOCKED);
+        if (isUnlockTransition) {
+            // When transitioning to unlocked, ensure the hotseat is fully visible from the
+            // beginning. The hotseat itself is animated by LauncherUnlockAnimationController.
             mIconAlignment.cancelAnimation();
             // updateValue ensures onIconAlignmentRatioChanged will be called if there is an actual
             // change in value
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index b142f9e..1e1769c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -585,9 +585,12 @@
         // already stashed Taskbar.
         boolean hotseatTopElement = mControllers.uiController.isHotseatIconOnTopWhenAligned()
                 || !hasAnyFlag(changedFlags, FLAG_IN_APP);
-        // If transitioning between locked/unlocked device, do not play a stash animation.
-        boolean unLockedTransition = hasAnyFlag(changedFlags, FLAG_STASHED_DEVICE_LOCKED);
-        boolean skipStashAnimation = !hotseatTopElement || unLockedTransition;
+        // If transitioning to unlocked device, do not play a stash animation.
+        // Keep isUnlockTransition in sync with its counterpart in
+        // TaskbarLauncherStateController#onStateChangeApplied.
+        boolean isUnlockTransition = hasAnyFlag(changedFlags, FLAG_STASHED_DEVICE_LOCKED)
+                && !hasAnyFlag(FLAG_STASHED_DEVICE_LOCKED);
+        boolean skipStashAnimation = !hotseatTopElement || isUnlockTransition;
 
         if (isTransientTaskbar) {
             createTransientAnimToIsStashed(mAnimator, isStashed, duration, animateBg, changedFlags,
@@ -909,8 +912,17 @@
                 SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE));
         updateStateForFlag(FLAG_STASHED_SYSUI,
                 hasAnyFlag(systemUiStateFlags, SYSUI_STATE_SCREEN_PINNING));
-        updateStateForFlag(FLAG_STASHED_DEVICE_LOCKED,
-                hasAnyFlag(systemUiStateFlags, MASK_ANY_SYSUI_LOCKED));
+
+        boolean isLocked = hasAnyFlag(systemUiStateFlags, MASK_ANY_SYSUI_LOCKED);
+        boolean wasLocked = hasAnyFlag(FLAG_STASHED_DEVICE_LOCKED);
+        updateStateForFlag(FLAG_STASHED_DEVICE_LOCKED, isLocked);
+
+        if (isLocked && !wasLocked && DisplayController.isTransientTaskbar(mActivity)) {
+            // Stash the transient taskbar when locking the device. This improves the transition
+            // to AoD (otherwise the taskbar stays a bit too long above the collapsing AoD scrim),
+            // and ensures the taskar state is reset when unlocking the device afterwards.
+            updateStateForFlag(FLAG_STASHED_IN_APP_AUTO, true);
+        }
 
         // Only update FLAG_STASHED_IN_APP_IME when system gesture is not in progress.
         mIsImeShowing = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_SHOWING);