diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
index 8a1b68ddda7d2..64cfb4bcd0585 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenBrightness.java
@@ -127,7 +127,9 @@ public class DozeScreenBrightness extends BroadcastReceiver implements DozeMachi
 
     @Override
     public void onScreenState(int state) {
-        if (mState == DozeMachine.State.FINISH && !mScreenOff
+        if (!mScreenOff
+                && (mState == DozeMachine.State.DOZE_AOD
+                     || mState == DozeMachine.State.DOZE_AOD_DOCKED)
                 && (state == Display.STATE_DOZE || state == Display.STATE_DOZE_SUSPEND)) {
             setLightSensorEnabled(true);
         } else {
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 5288f93dbc38b..3ef60274cd767 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenBrightnessTest.java
@@ -105,7 +105,6 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     public void testAod_usesLightSensor() {
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
-        mScreen.transitionTo(DOZE_AOD, FINISH);
         mScreen.onScreenState(Display.STATE_DOZE);
 
         mSensor.sendSensorEvent(3);
@@ -141,7 +140,6 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     public void testPausingAod_doesNotResetBrightness() throws Exception {
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
-        mScreen.transitionTo(DOZE_AOD, FINISH);
         mScreen.onScreenState(Display.STATE_DOZE);
 
         mSensor.sendSensorEvent(1);
@@ -210,7 +208,6 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     public void testNonPositiveBrightness_keepsPreviousBrightnessAndScrim() throws Exception {
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
-        mScreen.transitionTo(DOZE_AOD, FINISH);
         mScreen.onScreenState(Display.STATE_DOZE);
 
         mSensor.sendSensorEvent(1);
@@ -224,19 +221,17 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     public void pausingAod_unblanksAfterSensor() {
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
-        mScreen.transitionTo(DOZE_AOD, FINISH);
         mScreen.onScreenState(Display.STATE_DOZE);
 
         mSensor.sendSensorEvent(2);
 
-        mScreen.transitionTo(FINISH, DOZE_AOD_PAUSING);
+        mScreen.transitionTo(DOZE_AOD, DOZE_AOD_PAUSING);
         mScreen.transitionTo(DOZE_AOD_PAUSING, DOZE_AOD_PAUSED);
 
         mSensor.sendSensorEvent(0);
 
         reset(mDozeHost);
         mScreen.transitionTo(DOZE_AOD_PAUSED, DOZE_AOD);
-        mScreen.transitionTo(DOZE_AOD, FINISH);
         mScreen.onScreenState(Display.STATE_DOZE);
         mSensor.sendSensorEvent(2);
         verify(mDozeHost).setAodDimmingScrim(eq(0f));
@@ -246,11 +241,10 @@ public class DozeScreenBrightnessTest extends SysuiTestCase {
     public void pausingAod_unblanksIfSensorWasAlwaysReady() throws Exception {
         mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
         mScreen.transitionTo(INITIALIZED, DOZE_AOD);
-        mScreen.transitionTo(DOZE_AOD, FINISH);
         mScreen.onScreenState(Display.STATE_DOZE);
 
         mSensor.sendSensorEvent(2);
-        mScreen.transitionTo(FINISH, DOZE_AOD_PAUSING);
+        mScreen.transitionTo(DOZE_AOD, DOZE_AOD_PAUSING);
         mScreen.transitionTo(DOZE_AOD_PAUSING, DOZE_AOD_PAUSED);
 
         reset(mDozeHost);