diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
index 26d63b6e44a57..d27c39a893190 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
@@ -16,6 +16,10 @@
 
 package com.android.systemui.doze;
 
+import static android.os.PowerManager.GO_TO_SLEEP_REASON_TIMEOUT;
+
+import static com.android.systemui.keyguard.WakefulnessLifecycle.WAKEFULNESS_GOING_TO_SLEEP;
+
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
@@ -262,9 +266,11 @@ public class DozeScreenBrightness extends BroadcastReceiver implements DozeMachi
      * animation.
      */
     private int clampToDimBrightnessForScreenOff(int brightness) {
-        if (mUnlockedScreenOffAnimationController.isScreenOffAnimationPlaying()
-                && mWakefulnessLifecycle.getLastSleepReason()
-                == PowerManager.GO_TO_SLEEP_REASON_TIMEOUT) {
+        final boolean screenTurningOff =
+                mUnlockedScreenOffAnimationController.isScreenOffAnimationPlaying()
+                        || mWakefulnessLifecycle.getWakefulness() == WAKEFULNESS_GOING_TO_SLEEP;
+        if (screenTurningOff
+                && mWakefulnessLifecycle.getLastSleepReason() == GO_TO_SLEEP_REASON_TIMEOUT) {
             return Math.max(
                     PowerManager.BRIGHTNESS_OFF,
                     // Use the lower of either the dim brightness, or the current brightness reduced
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
index c021620b64e2b..5b472ba6557b2 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
@@ -475,7 +475,7 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     }
 
     @Test
-    public void transitionToDoze_duringScreenOff_afterTimeout_clampsToDim() {
+    public void transitionToDoze_duringUnlockedScreenOff_afterTimeout_clampsToDim() {
         when(mWakefulnessLifecycle.getLastSleepReason()).thenReturn(
                 PowerManager.GO_TO_SLEEP_REASON_TIMEOUT);
         when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(true);
@@ -490,7 +490,7 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     }
 
     @Test
-    public void transitionToDoze_duringScreenOff_notAfterTimeout_doesNotClampToDim() {
+    public void transitionToDoze_duringUnlockedScreenOff_notAfterTimeout_doesNotClampToDim() {
         when(mWakefulnessLifecycle.getLastSleepReason()).thenReturn(
                 PowerManager.GO_TO_SLEEP_REASON_POWER_BUTTON);
         when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(true);
@@ -505,7 +505,7 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     }
 
     @Test
-    public void transitionToDoze_duringScreenOff_afterTimeout_noScreenOff_doesNotClampToDim() {
+    public void transitionToDoze_duringUnlockedScreenOff_afterTimeout_noScreenOff_doesNotClampToDim() {
         when(mWakefulnessLifecycle.getLastSleepReason()).thenReturn(
                 PowerManager.GO_TO_SLEEP_REASON_TIMEOUT);
         when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(false);
@@ -518,6 +518,36 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
         assertEquals(mServiceFake.screenBrightness, DEFAULT_BRIGHTNESS);
     }
 
+    @Test
+    public void transitionToDoze_duringLockedScreenOff_afterTimeout_clampsToDim() {
+        when(mWakefulnessLifecycle.getLastSleepReason()).thenReturn(
+                PowerManager.GO_TO_SLEEP_REASON_TIMEOUT);
+        when(mWakefulnessLifecycle.getWakefulness()).thenReturn(
+                WakefulnessLifecycle.WAKEFULNESS_GOING_TO_SLEEP);
+        when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(false);
+        when(mUnlockedScreenOffAnimationController.isScreenOffAnimationPlaying()).thenReturn(false);
+
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+
+        assertTrue(mServiceFake.screenBrightness <= DIM_BRIGHTNESS);
+    }
+
+    @Test
+    public void transitionToDoze_duringLockedScreenOff_notAfterTimeout_doesNotClampToDim() {
+        when(mWakefulnessLifecycle.getLastSleepReason()).thenReturn(
+                PowerManager.GO_TO_SLEEP_REASON_POWER_BUTTON);
+        when(mWakefulnessLifecycle.getWakefulness()).thenReturn(
+                WakefulnessLifecycle.WAKEFULNESS_GOING_TO_SLEEP);
+        when(mDozeParameters.shouldControlUnlockedScreenOff()).thenReturn(false);
+        when(mUnlockedScreenOffAnimationController.isScreenOffAnimationPlaying()).thenReturn(false);
+
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mScreen.transitionTo(INITIALIZED, DOZE);
+
+        assertEquals(mServiceFake.screenBrightness, DEFAULT_BRIGHTNESS);
+    }
+
     private void waitForSensorManager() {
         mFakeExecutor.runAllReady();
     }