diff --git a/core/java/com/android/internal/util/StateMachine.java b/core/java/com/android/internal/util/StateMachine.java
index b0d45e1d1db9b..be10608df2a3f 100644
--- a/core/java/com/android/internal/util/StateMachine.java
+++ b/core/java/com/android/internal/util/StateMachine.java
@@ -1189,6 +1189,26 @@ public class StateMachine {
             return stateInfo;
         }
 
+        /**
+         * Remove a state from the state machine. Will not remove the state if it is currently
+         * active or if it has any children in the hierarchy.
+         * @param state the state to remove
+         */
+        private void removeState(State state) {
+            StateInfo stateInfo = mStateInfo.get(state);
+            if (stateInfo == null || stateInfo.active) {
+                return;
+            }
+            boolean isParent = mStateInfo.values().stream()
+                    .filter(si -> si.parentStateInfo == stateInfo)
+                    .findAny()
+                    .isPresent();
+            if (isParent) {
+                return;
+            }
+            mStateInfo.remove(state);
+        }
+
         /**
          * Constructor
          *
@@ -1336,6 +1356,14 @@ public class StateMachine {
         mSmHandler.addState(state, null);
     }
 
+    /**
+     * Removes a state from the state machine, unless it is currently active or if it has children.
+     * @param state state to remove
+     */
+    public final void removeState(State state) {
+        mSmHandler.removeState(state);
+    }
+
     /**
      * Set the initial state. This must be invoked before
      * and messages are sent to the state machine.