diff --git a/examples/simulator/01-simulator.py b/examples/simulator/01-simulator.py
index 464d2e5..4696ea7 100644
--- a/examples/simulator/01-simulator.py
+++ b/examples/simulator/01-simulator.py
@@ -14,7 +14,7 @@ def example_simulator_01():
     simulator = stormpy.simulator.create_simulator(model, seed=42)
     final_outcomes = dict()
     for n in range(1000):
-        for i in range(100):
+        while not simulator.is_done():
             observation = simulator.step()
         if observation not in final_outcomes:
             final_outcomes[observation] = 1
@@ -30,7 +30,7 @@ def example_simulator_01():
     simulator.set_observation_mode(stormpy.simulator.SimulatorObservationMode.PROGRAM_LEVEL)
     final_outcomes = dict()
     for n in range(1000):
-        for i in range(100):
+        while not simulator.is_done():
             observation = simulator.step()
         if observation not in final_outcomes:
             final_outcomes[observation] = 1
diff --git a/lib/stormpy/simulator.py b/lib/stormpy/simulator.py
index 638bbc1..3d39b82 100644
--- a/lib/stormpy/simulator.py
+++ b/lib/stormpy/simulator.py
@@ -16,15 +16,34 @@ class Simulator:
         self._observation_mode = SimulatorObservationMode.STATE_LEVEL
 
     def step(self, action=None):
+        """
+        Do a step taking the passed action.
+
+        :param action: The index of the action, for deterministic actions, action may be None.
+        :return: The observation (on state or program level).
+        """
         raise NotImplementedError("Abstract superclass")
 
     def restart(self):
+        """
+        Reset the simulator to the initial state
+        """
         raise NotImplementedError("Abstract superclass")
 
     def is_done(self):
+        """
+        Is the simulator in a sink state?
+
+        :return: Yes, if the simulator is in a sink state.
+        """
         return False
 
     def set_observation_mode(self, mode):
+        """
+
+        :param mode: STATE_LEVEL or PROGRAM_LEVEL
+        :type mode:
+        """
         if not isinstance(mode, SimulatorObservationMode):
             raise RuntimeError("Observation mode must be a SimulatorObservationMode")
         self._observation_mode = mode