diff --git a/sr_dqn_lab/train_script.py b/sr_dqn_lab/train_script.py
index cab2e6a..e03f302 100644
--- a/sr_dqn_lab/train_script.py
+++ b/sr_dqn_lab/train_script.py
@@ -1,6 +1,7 @@
 import os
 import gymnasium as gym
 import numpy as np
+import minigrid
 import argparse
 import matplotlib.pyplot as plt
 from stable_baselines3 import DQN