diff --git a/packages/SystemUI/src/com/android/systemui/doze/AlwaysOnDisplayPolicy.java b/packages/SystemUI/src/com/android/systemui/doze/AlwaysOnDisplayPolicy.java
index 8515bf2e2968a..4fea45c39d5df 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/AlwaysOnDisplayPolicy.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/AlwaysOnDisplayPolicy.java
@@ -116,6 +116,7 @@ public class AlwaysOnDisplayPolicy {
     private SettingsObserver mSettingsObserver;
 
     public AlwaysOnDisplayPolicy(Context context) {
+        context = context.getApplicationContext();
         mContext = context;
         mParser = new KeyValueListParser(',');
         mSettingsObserver = new SettingsObserver(context.getMainThreadHandler());
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java b/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
index 5bf62f638235e..e9d78d9b4fdf5 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
@@ -61,13 +61,14 @@ public class DozeFactory {
 
         DozeMachine machine = new DozeMachine(wrappedService, config, wakeLock);
         machine.setParts(new DozeMachine.Part[]{
-                new DozePauser(handler, machine, alarmManager, new AlwaysOnDisplayPolicy(context)),
+                new DozePauser(handler, machine, alarmManager, params.getPolicy()),
                 new DozeFalsingManagerAdapter(FalsingManager.getInstance(context)),
                 createDozeTriggers(context, sensorManager, host, alarmManager, config, params,
                         handler, wakeLock, machine),
                 createDozeUi(context, host, wakeLock, machine, handler, alarmManager, params),
                 new DozeScreenState(wrappedService, handler, params, wakeLock),
-                createDozeScreenBrightness(context, wrappedService, sensorManager, host, handler),
+                createDozeScreenBrightness(context, wrappedService, sensorManager, host, params,
+                        handler),
                 new DozeWallpaperState(context, params)
         });
 
@@ -76,11 +77,11 @@ public class DozeFactory {
 
     private DozeMachine.Part createDozeScreenBrightness(Context context,
             DozeMachine.Service service, SensorManager sensorManager, DozeHost host,
-            Handler handler) {
+            DozeParameters params, Handler handler) {
         Sensor sensor = DozeSensors.findSensorWithType(sensorManager,
                 context.getString(R.string.doze_brightness_sensor_type));
         return new DozeScreenBrightness(context, service, sensorManager, sensor, host, handler,
-                new AlwaysOnDisplayPolicy(context));
+                params.getPolicy());
     }
 
     private DozeTriggers createDozeTriggers(Context context, SensorManager sensorManager,
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeService.java b/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
index aa264190f8007..73393047cc455 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeService.java
@@ -62,6 +62,7 @@ public class DozeService extends DreamService
     public void onDestroy() {
         Dependency.get(PluginManager.class).removePluginListener(this);
         super.onDestroy();
+        mDozeMachine = null;
     }
 
     @Override
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
index f7a258a2c959d..be3e322f598fd 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
@@ -85,7 +85,7 @@ public class DozeTriggers implements DozeMachine.Part {
         mAllowPulseTriggers = allowPulseTriggers;
         mDozeSensors = new DozeSensors(context, alarmManager, mSensorManager, dozeParameters,
                 config, wakeLock, this::onSensor, this::onProximityFar,
-                new AlwaysOnDisplayPolicy(context));
+                dozeParameters.getPolicy());
         mUiModeManager = mContext.getSystemService(UiModeManager.class);
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
index 07b79a209b823..d2d9c4c9613bc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
@@ -59,9 +59,9 @@ public class DozeParameters implements TunerService.Tunable {
 
     @VisibleForTesting
     protected DozeParameters(Context context) {
-        mContext = context;
+        mContext = context.getApplicationContext();
         mAmbientDisplayConfiguration = new AmbientDisplayConfiguration(mContext);
-        mAlwaysOnPolicy = new AlwaysOnDisplayPolicy(context);
+        mAlwaysOnPolicy = new AlwaysOnDisplayPolicy(mContext);
 
         mControlScreenOffAnimation = !getDisplayNeedsBlanking();
         mPowerManager = mContext.getSystemService(PowerManager.class);
@@ -243,6 +243,10 @@ public class DozeParameters implements TunerService.Tunable {
         mDozeAlwaysOn = mAmbientDisplayConfiguration.alwaysOnEnabled(UserHandle.USER_CURRENT);
     }
 
+    public AlwaysOnDisplayPolicy getPolicy() {
+        return mAlwaysOnPolicy;
+    }
+
     /**
      * Parses a spec of the form `1,2,3,!5,*`. The resulting object will match numbers that are
      * listed, will not match numbers that are listed with a ! prefix, and will match / not match