diff --git a/main.py b/main.py
index af5a561..cf0ea0f 100755
--- a/main.py
+++ b/main.py
@@ -2,7 +2,7 @@
 import argparse
 
 from loguru import logger
-from utils import BASE_PATH, CONFIG
+from utils import BASE_PATH, CONFIG, GameMode
 
 
 def pos_int(string: str) -> int:
@@ -66,10 +66,11 @@ def main(args: argparse.ArgumentParser) -> None:
 
     if args.train is not None:
         ai.log.debug("Training the AI")
-        ai.train(*args.train)
+        # ai.train(*args.train)
+        game.Main(GameMode.AI_TRAINING).run()
     else:
         game.log.debug("Running the game")
-        game.Main().run()
+        game.Main(GameMode.PLAYER).run()
 
 
 if __name__ == "__main__":
diff --git a/src/game/main.py b/src/game/main.py
index 96f873b..3dd5acc 100644
--- a/src/game/main.py
+++ b/src/game/main.py
@@ -1,7 +1,7 @@
 import sys
 
 import pygame
-from utils import CONFIG, Figure
+from utils import CONFIG, Figure, GameMode
 
 from .game import Game
 from .log import log
@@ -25,8 +25,9 @@ class Main:
         music: Pygame music that plays in the background.
     """
 
-    def __init__(self) -> None:
+    def __init__(self, mode: GameMode) -> None:
         log.info("Initializing the game")
+        self.game_mode = mode
         self._initialize_pygeme()
         self._initialize_game_components()
         self._start_background_music()
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index e93aab5..6870edc 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -1,5 +1,5 @@
 from .config import CONFIG
-from .enum import Direction, Field, Rotation
+from .enum import Direction, Field, GameMode, Rotation
 from .figure import Figure, FigureConfig
 from .log import log
 from .path import BASE_PATH
@@ -15,4 +15,5 @@ __all__ = [
     "Direction",
     "Field",
     "Rotation",
+    "GameMode",
 ]
diff --git a/src/utils/enum.py b/src/utils/enum.py
index c9bad06..438346a 100644
--- a/src/utils/enum.py
+++ b/src/utils/enum.py
@@ -1,4 +1,4 @@
-from enum import Enum
+from enum import Enum, auto
 
 
 class Direction(Enum):
@@ -16,3 +16,9 @@ class Rotation(Enum):
 class Field(Enum):
     EMPTY = None
     FILLED = "Block"
+
+
+class GameMode(Enum):
+    PLAYER = auto()
+    AI_PLAYING = auto()
+    AI_TRAINING = auto()