diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonDrawable.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonDrawable.java
index 8fcaa67e76149..23d03a4b225a0 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonDrawable.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonDrawable.java
@@ -80,6 +80,22 @@ public class KeyButtonDrawable extends Drawable {
     private final Paint mShadowPaint = new Paint(Paint.ANTI_ALIAS_FLAG | Paint.FILTER_BITMAP_FLAG);
     private final ShadowDrawableState mState;
     private AnimatedVectorDrawable mAnimatedDrawable;
+    private final Callback mAnimatedDrawableCallback = new Callback() {
+        @Override
+        public void invalidateDrawable(@NonNull Drawable who) {
+            invalidateSelf();
+        }
+
+        @Override
+        public void scheduleDrawable(@NonNull Drawable who, @NonNull Runnable what, long when) {
+            scheduleSelf(what, when);
+        }
+
+        @Override
+        public void unscheduleDrawable(@NonNull Drawable who, @NonNull Runnable what) {
+            unscheduleSelf(what);
+        }
+    };
 
     public KeyButtonDrawable(Drawable d, @ColorInt int lightColor, @ColorInt int darkColor,
             boolean horizontalFlip, Color ovalBackgroundColor) {
@@ -97,6 +113,7 @@ public class KeyButtonDrawable extends Drawable {
         }
         if (canAnimate()) {
             mAnimatedDrawable = (AnimatedVectorDrawable) mState.mChildState.newDrawable().mutate();
+            mAnimatedDrawable.setCallback(mAnimatedDrawableCallback);
             setDrawableBounds(mAnimatedDrawable);
         }
     }