diff --git a/services/core/java/com/android/server/display/LogicalDisplayMapper.java b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
index 375e51c..1f7232a 100644
--- a/services/core/java/com/android/server/display/LogicalDisplayMapper.java
+++ b/services/core/java/com/android/server/display/LogicalDisplayMapper.java
@@ -159,6 +159,7 @@
     private Layout mCurrentLayout = null;
     private int mDeviceState = DeviceStateManager.INVALID_DEVICE_STATE;
     private int mPendingDeviceState = DeviceStateManager.INVALID_DEVICE_STATE;
+    private int mDeviceStateToBeAppliedAfterBoot = DeviceStateManager.INVALID_DEVICE_STATE;
     private boolean mBootCompleted = false;
     private boolean mInteractive;
 
@@ -353,6 +354,12 @@
         ipw.println("mDeviceStatesOnWhichToWakeUp=" + mDeviceStatesOnWhichToWakeUp);
         ipw.println("mDeviceStatesOnWhichToSleep=" + mDeviceStatesOnWhichToSleep);
         ipw.println("mInteractive=" + mInteractive);
+        ipw.println("mBootCompleted=" + mBootCompleted);
+
+        ipw.println();
+        ipw.println("mDeviceState=" + mDeviceState);
+        ipw.println("mPendingDeviceState=" + mPendingDeviceState);
+        ipw.println("mDeviceStateToBeAppliedAfterBoot=" + mDeviceStateToBeAppliedAfterBoot);
 
         final int logicalDisplayCount = mLogicalDisplays.size();
         ipw.println();
@@ -370,10 +377,6 @@
     }
 
     void setDeviceStateLocked(int state, boolean isOverrideActive) {
-        Slog.i(TAG, "Requesting Transition to state: " + state + ", from state=" + mDeviceState
-                + ", interactive=" + mInteractive + ", mBootCompleted=" + mBootCompleted);
-        mPendingDeviceState = state;
-
         if (!mBootCompleted) {
             // The boot animation might still be in progress, we do not want to switch states now
             // as the boot animation would end up with an incorrect size.
@@ -381,14 +384,19 @@
                 Slog.d(TAG, "Postponing transition to state: " + mPendingDeviceState
                         + " until boot is completed");
             }
+            mDeviceStateToBeAppliedAfterBoot = state;
             return;
         }
 
+        Slog.i(TAG, "Requesting Transition to state: " + state + ", from state=" + mDeviceState
+                + ", interactive=" + mInteractive + ", mBootCompleted=" + mBootCompleted);
         // As part of a state transition, we may need to turn off some displays temporarily so that
         // the transition is smooth. Plus, on some devices, only one internal displays can be
         // on at a time. We use LogicalDisplay.setIsInTransition to mark a display that needs to be
         // temporarily turned off.
         resetLayoutLocked(mDeviceState, state, /* isStateChangeStarting= */ true);
+        mPendingDeviceState = state;
+        mDeviceStateToBeAppliedAfterBoot = DeviceStateManager.INVALID_DEVICE_STATE;
         final boolean wakeDevice = shouldDeviceBeWoken(mPendingDeviceState, mDeviceState,
                 mInteractive, mBootCompleted);
         final boolean sleepDevice = shouldDeviceBePutToSleep(mPendingDeviceState, mDeviceState,
@@ -435,8 +443,9 @@
     void onBootCompleted() {
         synchronized (mSyncRoot) {
             mBootCompleted = true;
-            if (mPendingDeviceState != DeviceStateManager.INVALID_DEVICE_STATE) {
-                setDeviceStateLocked(mPendingDeviceState, /* isOverrideActive= */ false);
+            if (mDeviceStateToBeAppliedAfterBoot != DeviceStateManager.INVALID_DEVICE_STATE) {
+                setDeviceStateLocked(mDeviceStateToBeAppliedAfterBoot,
+                        /* isOverrideActive= */ false);
             }
         }
     }
@@ -491,7 +500,8 @@
     @VisibleForTesting
     boolean shouldDeviceBePutToSleep(int pendingState, int currentState, boolean isOverrideActive,
             boolean isInteractive, boolean isBootCompleted) {
-        return mDeviceStatesOnWhichToSleep.get(pendingState)
+        return currentState != DeviceStateManager.INVALID_DEVICE_STATE
+                && mDeviceStatesOnWhichToSleep.get(pendingState)
                 && !mDeviceStatesOnWhichToSleep.get(currentState)
                 && !isOverrideActive
                 && isInteractive && isBootCompleted;
diff --git a/services/tests/servicestests/src/com/android/server/display/LogicalDisplayMapperTest.java b/services/tests/servicestests/src/com/android/server/display/LogicalDisplayMapperTest.java
index 6790ad9..650eef0 100644
--- a/services/tests/servicestests/src/com/android/server/display/LogicalDisplayMapperTest.java
+++ b/services/tests/servicestests/src/com/android/server/display/LogicalDisplayMapperTest.java
@@ -591,6 +591,7 @@
         // 2) Mark the displays as STATE_OFF so that it can continue with transition
         // 3) Send DISPLAY_DEVICE_EVENT_CHANGE to inform the mapper of the new display state
         // 4) Dispatch handler events.
+        mLogicalDisplayMapper.onBootCompleted();
         mLogicalDisplayMapper.setDeviceStateLocked(0, false);
         mDisplayDeviceRepo.onDisplayDeviceEvent(device3, DISPLAY_DEVICE_EVENT_CHANGED);
         mLooper.moveTimeForward(1000);