diff --git a/src/vivarium/framework/state_machine.py b/src/vivarium/framework/state_machine.py
index 7a2ab9a4..c44d7b6f 100644
--- a/src/vivarium/framework/state_machine.py
+++ b/src/vivarium/framework/state_machine.py
@@ -563,6 +563,9 @@ def on_initialize_simulants(self, pop_data: SimulantData) -> None:
     def on_time_step(self, event: Event) -> None:
         self.transition(event.index, event.time)
 
+    def on_time_step_cleanup(self, event: Event) -> None:
+        self.cleanup(event.index, event.time)
+
     ##################
     # Public methods #
     ##################