diff --git a/torch_training/training_navigation.py b/torch_training/training_navigation.py
index fe4905f91cb84964ca27f8dded888a19a2146e6c..e03c5d15fec02b11659d09b340dce17c82971c0d 100644
--- a/torch_training/training_navigation.py
+++ b/torch_training/training_navigation.py
@@ -42,10 +42,11 @@ def main(argv):
 
 
     # Use a the malfunction generator to break agents from time to time
-    stochastic_data = {'malfunction_rate': 8000,  # Rate of malfunction occurence of single agent
-                       'min_duration': 15,  # Minimal duration of malfunction
-                       'max_duration': 50  # Max duration of malfunction
-                       }
+    stochastic_data = MalfunctionParameters(malfunction_rate=10000,  # Rate of malfunction occurence
+                                            min_duration=15,  # Minimal duration of malfunction
+                                            max_duration=50  # Max duration of malfunction
+                                            )
+
 
     # Custom observation builder
     TreeObservation = TreeObsForRailEnv(max_depth=2)
@@ -69,7 +70,8 @@ def main(argv):
                   malfunction_generator_and_process_data=malfunction_from_params(stochastic_data),
                   # Malfunction data generator
                   obs_builder_object=TreeObservation)
-
+    # Reset env
+    env.reset(True,True)
     # After training we want to render the results so we also load a renderer
     env_renderer = RenderTool(env, gl="PILSVG", )
     # Given the depth of the tree observation and the number of features per node we get the following state_size