diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py
index 8b771cd7..7a183588 100644
--- a/src/primaite/simulator/core.py
+++ b/src/primaite/simulator/core.py
@@ -177,7 +177,7 @@ class SimComponent(BaseModel):
         """
         pass
 
-    def reset_component_for_episode(self):
+    def reset_component_for_episode(self, episode: int):
         """
         Reset this component to its original state for a new episode.