diff --git a/resources/graphs/state_machine/graph.py b/resources/graphs/state_machine/graph.py
index 4e39cc5c0ed3e32b1b7c7d7c1628524e3a376383..e5e22b7ec4dc91979622737cbe0eaaaa59018c20 100644
--- a/resources/graphs/state_machine/graph.py
+++ b/resources/graphs/state_machine/graph.py
@@ -353,3 +353,22 @@ class StateMachineGraph(DirectedGraph):
         self._edges[connecting_edge.get_name()] = connecting_edge
 
         return connecting_edge
+
+    def remove_node(self, node_indentifier):
+        """
+        Removes node from graph.
+        :param node_indentifier: Node or name of node to remove.
+        :return: Removed node.
+        """
+        node = self.get_node(node_indentifier)
+
+        # Remove initial state, if present.
+        if node.is_initial:
+            self._initial_states.pop(node.get_name())
+
+        # Remove final state, if present.
+        if node.is_final:
+            self._final_states.pop(node.get_name())
+
+        # Call parent implementation.
+        return super().remove_node(node_indentifier)