diff --git a/torch_training/render_agent_behavior.py b/torch_training/render_agent_behavior.py
index 82706a4d2e22df7b11f03207d0b7d6aac891a89c..bfe5ba030d265fa1f765946d418dbd9e526558b8 100644
--- a/torch_training/render_agent_behavior.py
+++ b/torch_training/render_agent_behavior.py
@@ -3,7 +3,7 @@ from collections import deque
 
 import numpy as np
 import torch
-from flatland.envs.malfunction_generators import malfunction_from_params
+from flatland.envs.malfunction_generators import malfunction_from_params, MalfunctionParameters
 from flatland.envs.observations import TreeObsForRailEnv
 from flatland.envs.rail_env import RailEnv
 from flatland.envs.rail_generators import sparse_rail_generator