ai-snake-lab 0.1.0__py3-none-any.whl → 0.4.4__py3-none-any.whl

ai_snake_lab/AISim.py

@@ -10,44 +10,74 @@ AISim.py
 
 import threading
 import time
-import sys
+import sys, os
+from datetime import datetime, timedelta
 
 from textual.app import App, ComposeResult
-from textual.widgets import Label, Input, Button
+from textual.widgets import Label, Input, Button, Static
 from textual.containers import Vertical, Horizontal
 from textual.reactive import var
+from textual.theme import Theme
 
-from constants.DDef import DDef
-from constants.DEpsilon import DEpsilon
-from constants.DFields import DField
-from constants.DFile import DFile
-from constants.DLayout import DLayout
-from constants.DLabels import DLabel
-from constants.DReplayMemory import MEM_TYPE
+from ai_snake_lab.constants.DDef import DDef
+from ai_snake_lab.constants.DEpsilon import DEpsilon
+from ai_snake_lab.constants.DFields import DField
+from ai_snake_lab.constants.DFile import DFile
+from ai_snake_lab.constants.DLayout import DLayout
+from ai_snake_lab.constants.DLabels import DLabel
+from ai_snake_lab.constants.DReplayMemory import MEM_TYPE
+from ai_snake_lab.constants.DDir import DDir
+from ai_snake_lab.constants.DDb4EPlot import Plot
 
-from ai.AIAgent import AIAgent
-from ai.EpsilonAlgo import EpsilonAlgo
-from game.GameBoard import GameBoard
-from game.SnakeGame import SnakeGame
+
+from ai_snake_lab.ai.AIAgent import AIAgent
+from ai_snake_lab.ai.EpsilonAlgo import EpsilonAlgo
+from ai_snake_lab.game.GameBoard import GameBoard
+from ai_snake_lab.game.SnakeGame import SnakeGame
+from ai_snake_lab.ui.Db4EPlot import Db4EPlot
 
 RANDOM_SEED = 1970
 
+snake_lab_theme = Theme(
+    name="db4e",
+    primary="#88C0D0",
+    secondary="#1f6a83ff",
+    accent="#B48EAD",
+    foreground="#31b8e6",
+    background="black",
+    success="#A3BE8C",
+    warning="#EBCB8B",
+    error="#BF616A",
+    surface="black",
+    panel="#000000",
+    dark=True,
+    variables={
+        "block-cursor-text-style": "none",
+        "footer-key-foreground": "#88C0D0",
+        "input-selection-background": "#81a1c1 35%",
+    },
+)
+
 
 class AISim(App):
     """A Textual app that has an AI Agent playing the Snake Game."""
 
     TITLE = DDef.APP_TITLE
-    CSS_PATH = DFile.CSS_PATH
+    CSS_PATH = os.path.join(DDir.UTILS, DFile.CSS_FILE)
 
     ## Runtime values
     # Current epsilon value (degrades in real-time)
     cur_epsilon_widget = Label("N/A", id=DLayout.CUR_EPSILON)
     # Current memory type
     cur_mem_type_widget = Label("N/A", id=DLayout.CUR_MEM_TYPE)
-    # Number of stored memories
-    cur_num_memories_widget = Label("N/A", id=DLayout.NUM_MEMORIES)
-    # Runtime move delay value
+    # Current model type
+    cur_model_type_widget = Label("N/A", id=DLayout.CUR_MODEL_TYPE)
+    # Time delay between moves
     cur_move_delay = DDef.MOVE_DELAY
+    # Number of stored games in the ReplayMemory
+    cur_num_games_widget = Label("N/A", id=DLayout.NUM_GAMES)
+    # Elapsed time
+    cur_runtime_widget = Label("N/A", id=DLayout.RUNTIME)
 
     # Intial Settings for Epsilon
     initial_epsilon_input = Input(
@@ -77,23 +107,42 @@ class AISim(App):
 
     # Buttons
     pause_button = Button(label=DLabel.PAUSE, id=DLayout.BUTTON_PAUSE, compact=True)
+    restart_button = Button(
+        label=DLabel.RESTART, id=DLayout.BUTTON_RESTART, compact=True
+    )
     start_button = Button(label=DLabel.START, id=DLayout.BUTTON_START, compact=True)
     quit_button = Button(label=DLabel.QUIT, id=DLayout.BUTTON_QUIT, compact=True)
-    reset_button = Button(label=DLabel.RESET, id=DLayout.BUTTON_RESET, compact=True)
+    defaults_button = Button(
+        label=DLabel.DEFAULTS, id=DLayout.BUTTON_DEFAULTS, compact=True
+    )
     update_button = Button(label=DLabel.UPDATE, id=DLayout.BUTTON_UPDATE, compact=True)
 
+    # A dictionary to hold runtime statistics
+    stats = {
+        DField.GAME_SCORE: {
+            DField.GAME_NUM: [],
+            DField.GAME_SCORE: [],
+        }
+    }
+
+    game_score_plot = Db4EPlot(
+        title=DLabel.GAME_SCORE, id=DLayout.GAME_SCORE_PLOT, thin_method=Plot.SLIDING
+    )
+
     def __init__(self) -> None:
         super().__init__()
         self.game_board = GameBoard(20, id=DLayout.GAME_BOARD)
         self.snake_game = SnakeGame(game_board=self.game_board, id=DLayout.GAME_BOARD)
         self.epsilon_algo = EpsilonAlgo(seed=RANDOM_SEED)
         self.agent = AIAgent(self.epsilon_algo, seed=RANDOM_SEED)
-        self.running = False
-
-        self.score = Label("Game: 0, Highscore: 0, Score: 0")
+        self.cur_state = DField.STOPPED
+        self.game_score_plot._x_label = DLabel.GAME_NUM
+        self.game_score_plot._y_label = DLabel.GAME_SCORE
 
         # Setup the simulator in a background thread
         self.stop_event = threading.Event()
+        self.pause_event = threading.Event()
+        self.running = DField.STOPPED
         self.simulator_thread = threading.Thread(target=self.start_sim, daemon=True)
 
     async def action_quit(self) -> None:
@@ -105,70 +154,95 @@ class AISim(App):
 
     def compose(self) -> ComposeResult:
         """Create child widgets for the app."""
+
+        # Title bar
         yield Label(DDef.APP_TITLE, id=DLayout.TITLE)
-        yield Horizontal(
-            Vertical(
-                Vertical(
-                    Horizontal(
-                        Label(
-                            f"{DLabel.EPSILON_INITIAL} : ",
-                            classes=DLayout.LABEL_SETTINGS,
-                        ),
-                        self.initial_epsilon_input,
-                    ),
-                    Horizontal(
-                        Label(
-                            f"{DLabel.EPSILON_DECAY}   : ",
-                            classes=DLayout.LABEL_SETTINGS,
-                        ),
-                        self.epsilon_decay_input,
-                    ),
-                    Horizontal(
-                        Label(
-                            f"{DLabel.EPSILON_MIN} : ", classes=DLayout.LABEL_SETTINGS
-                        ),
-                        self.epsilon_min_input,
-                    ),
-                    Horizontal(
-                        Label(
-                            f"{DLabel.MOVE_DELAY}      : ",
-                            classes=DLayout.LABEL_SETTINGS,
-                        ),
-                        self.move_delay_input,
-                    ),
-                    id=DLayout.SETTINGS_BOX,
+
+        # Configuration Settings
+        yield Vertical(
+            Horizontal(
+                Label(
+                    f"{DLabel.EPSILON_INITIAL}",
+                    classes=DLayout.LABEL_SETTINGS,
                 ),
-                Vertical(
-                    Horizontal(
-                        self.start_button,
-                        self.reset_button,
-                        self.update_button,
-                        self.quit_button,
-                    ),
-                    id=DLayout.BUTTON_ROW,
+                self.initial_epsilon_input,
+            ),
+            Horizontal(
+                Label(
+                    f"{DLabel.EPSILON_DECAY}",
+                    classes=DLayout.LABEL_SETTINGS,
                 ),
+                self.epsilon_decay_input,
             ),
-            Vertical(
-                self.game_board,
-                id=DLayout.GAME_BOX,
+            Horizontal(
+                Label(f"{DLabel.EPSILON_MIN}", classes=DLayout.LABEL_SETTINGS),
+                self.epsilon_min_input,
             ),
-            Vertical(
-                Horizontal(
-                    Label(f"{DLabel.EPSILON} : ", classes=DLayout.LABEL),
-                    self.cur_epsilon_widget,
+            Horizontal(
+                Label(
+                    f"{DLabel.MOVE_DELAY}",
+                    classes=DLayout.LABEL_SETTINGS,
                 ),
-                Horizontal(
-                    Label(f"{DLabel.MEM_TYPE} : ", classes=DLayout.LABEL),
-                    self.cur_mem_type_widget,
-                ),
-                Horizontal(
-                    Label(f"{DLabel.MEMORIES} : ", classes=DLayout.LABEL),
-                    self.cur_num_memories_widget,
-                ),
-                id=DLayout.RUNTIME_BOX,
+                self.move_delay_input,
+            ),
+            id=DLayout.SETTINGS_BOX,
+        )
+
+        # The Snake Game
+        yield Vertical(
+            self.game_board,
+            id=DLayout.GAME_BOX,
+        )
+
+        # Runtime values
+        yield Vertical(
+            Horizontal(
+                Label(f"{DLabel.EPSILON}", classes=DLayout.LABEL),
+                self.cur_epsilon_widget,
+            ),
+            Horizontal(
+                Label(f"{DLabel.MEM_TYPE}", classes=DLayout.LABEL),
+                self.cur_mem_type_widget,
+            ),
+            Horizontal(
+                Label(f"{DLabel.STORED_GAMES}", classes=DLayout.LABEL),
+                self.cur_num_games_widget,
+            ),
+            Horizontal(
+                Label(f"{DLabel.MODEL_TYPE}", classes=DLayout.LABEL),
+                self.cur_model_type_widget,
+            ),
+            Horizontal(
+                Label(f"{DLabel.RUNTIME}", classes=DLayout.LABEL),
+                self.cur_runtime_widget,
             ),
+            id=DLayout.RUNTIME_BOX,
         )
 
+        # Buttons
+        yield Vertical(
+            Horizontal(
+                self.start_button,
+                self.pause_button,
+                self.quit_button,
+                classes=DLayout.BUTTON_ROW,
+            ),
+            Horizontal(
+                self.defaults_button,
+                self.update_button,
+                self.restart_button,
+                classes=DLayout.BUTTON_ROW,
+            ),
+        )
+
+        # Empty fillers
+        yield Static(id=DLayout.FILLER_1)
+        yield Static(id=DLayout.FILLER_2)
+        yield Static(id=DLayout.FILLER_3)
+
+        # The game score plot
+        yield self.game_score_plot
+
     def on_mount(self):
         self.initial_epsilon_input.value = str(DEpsilon.EPSILON_INITIAL)
         self.epsilon_decay_input.value = str(DEpsilon.EPSILON_DECAY)
@@ -177,16 +251,21 @@ class AISim(App):
         settings_box = self.query_one(f"#{DLayout.SETTINGS_BOX}", Vertical)
         settings_box.border_title = DLabel.SETTINGS
         runtime_box = self.query_one(f"#{DLayout.RUNTIME_BOX}", Vertical)
-        runtime_box.border_title = DLabel.RUNTIME
+        runtime_box.border_title = DLabel.RUNTIME_VALUES
         self.cur_mem_type_widget.update(
             MEM_TYPE.MEM_TYPE_TABLE[self.agent.memory.mem_type()]
         )
-        self.cur_num_memories_widget.update(str(self.agent.memory.get_num_memories()))
+        self.cur_num_games_widget.update(str(self.agent.memory.get_num_games()))
         # Initial state is that the app is stopped
         self.add_class(DField.STOPPED)
+        # Register the theme
+        self.register_theme(snake_lab_theme)
+
+        # Set the app's theme
+        self.theme = "db4e"
 
     def on_quit(self):
-        if self.running == True:
+        if self.running == DField.RUNNING:
             self.stop_event.set()
             if self.simulator_thread.is_alive():
                 self.simulator_thread.join()
@@ -194,22 +273,69 @@ class AISim(App):
 
     def on_button_pressed(self, event: Button.Pressed) -> None:
         button_id = event.button.id
+
+        # Pause button was pressed
+        if button_id == DLayout.BUTTON_PAUSE:
+            self.pause_event.set()
+            self.running = DField.PAUSED
+            self.remove_class(DField.RUNNING)
+            self.add_class(DField.PAUSED)
+            self.cur_move_delay = float(self.move_delay_input.value)
+            self.cur_model_type_widget.update(self.agent.model_type())
+
+        # Restart button was pressed
+        elif button_id == DLayout.BUTTON_RESTART:
+            self.running = DField.STOPPED
+            self.add_class(DField.STOPPED)
+            self.remove_class(DField.PAUSED)
+
+            # Signal thread to stop
+            self.stop_event.set()
+            # Unpause so we can exit cleanly
+            self.pause_event.clear()
+            # Join the old thread
+            if self.simulator_thread.is_alive():
+                self.simulator_thread.join(timeout=2)
+
+            # Reset the game and the UI
+            self.snake_game.reset()
+            score = 0
+            highscore = 0
+            self.epoch = 1
+            game_box = self.query_one(f"#{DLayout.GAME_BOX}", Vertical)
+            game_box.border_title = ""
+            game_box.border_subtitle = ""
+
+            # Recreate events and get a new thread
+            self.stop_event = threading.Event()
+            self.pause_event = threading.Event()
+            self.simulator_thread = threading.Thread(target=self.start_sim, daemon=True)
+
         # Start button was pressed
-        if button_id == DLayout.BUTTON_START:
-            self.start_thread()
-            self.running = True
+        elif button_id == DLayout.BUTTON_START:
+            if self.running == DField.STOPPED:
+                self.start_thread()
+            elif self.running == DField.PAUSED:
+                self.pause_event.clear()
+            self.pause_event.clear()
+            self.running = DField.RUNNING
             self.add_class(DField.RUNNING)
             self.remove_class(DField.STOPPED)
+            self.remove_class(DField.PAUSED)
             self.cur_move_delay = float(self.move_delay_input.value)
+            self.cur_model_type_widget.update(self.agent.model_type())
+
         # Reset button was pressed
-        elif button_id == DLayout.BUTTON_RESET:
+        elif button_id == DLayout.BUTTON_DEFAULTS:
             self.initial_epsilon_input.value = str(DEpsilon.EPSILON_INITIAL)
             self.epsilon_decay_input.value = str(DEpsilon.EPSILON_DECAY)
             self.epsilon_min_input.value = str(DEpsilon.EPSILON_MIN)
             self.move_delay_input.value = str(DDef.MOVE_DELAY)
+
         # Quit button was pressed
         elif button_id == DLayout.BUTTON_QUIT:
             self.on_quit()
+
         # Update button was pressed
         elif button_id == DLayout.BUTTON_UPDATE:
             self.cur_move_delay = float(self.move_delay_input.value)
@@ -224,8 +350,13 @@ class AISim(App):
         self.epoch = 1
         game_box = self.query_one(f"#{DLayout.GAME_BOX}", Vertical)
         game_box.border_title = f"{DLabel.GAME} #{self.epoch}"
+        start_time = datetime.now()
 
         while not self.stop_event.is_set():
+            if self.pause_event.is_set():
+                self.pause_event.wait()
+                time.sleep(0.2)
+                continue
             # The actual training loop...
             old_state = game_board.get_state()
             move = agent.get_move(old_state)
@@ -261,14 +392,40 @@ class AISim(App):
                 else:
                     self.cur_epsilon_widget.update(str(round(cur_epsilon, 4)))
                 # Update the number of stored memories
-                self.cur_num_memories_widget.update(
-                    str(self.agent.memory.get_num_memories())
-                )
+                self.cur_num_games_widget.update(str(self.agent.memory.get_num_games()))
+                # Update the stats object
+                self.stats[DField.GAME_SCORE][DField.GAME_NUM].append(self.epoch)
+                self.stats[DField.GAME_SCORE][DField.GAME_SCORE].append(score)
+                # Update the plot object
+                self.game_score_plot.add_data(self.epoch, score)
+                self.game_score_plot.db4e_plot()
+                # Update the runtime widget
+                elapsed_secs = (datetime.now() - start_time).total_seconds()
+                runtime = minutes_to_uptime(elapsed_secs)
+                self.cur_runtime_widget.update(runtime)
 
     def start_thread(self):
         self.simulator_thread.start()
 
 
+def minutes_to_uptime(seconds: int):
+    # Return a string like:
+    # 0h 0m 45s
+    # 1d 7h 32m
+    days, minutes = divmod(int(seconds), 86400)
+    hours, minutes = divmod(minutes, 3600)
+    minutes, seconds = divmod(minutes, 60)
+
+    if days > 0:
+        return f"{days}d {hours}h {minutes}m"
+    elif hours > 0:
+        return f"{hours}h {minutes}m"
+    elif minutes > 0:
+        return f"{minutes}m {seconds}s"
+    else:
+        return f"{seconds}s"
+
+
 if __name__ == "__main__":
     app = AISim()
     app.run()

ai_snake_lab/ai/AIAgent.py

@@ -9,13 +9,14 @@ ai/Agent.py
 """
 
 import torch
-from ai.EpsilonAlgo import EpsilonAlgo
-from ai.ReplayMemory import ReplayMemory
-from ai.AITrainer import AITrainer
-from ai.models.ModelL import ModelL
-from ai.models.ModelRNN import ModelRNN
+from ai_snake_lab.ai.EpsilonAlgo import EpsilonAlgo
+from ai_snake_lab.ai.ReplayMemory import ReplayMemory
+from ai_snake_lab.ai.AITrainer import AITrainer
+from ai_snake_lab.ai.models.ModelL import ModelL
+from ai_snake_lab.ai.models.ModelRNN import ModelRNN
 
-from constants.DReplayMemory import MEM_TYPE
+from ai_snake_lab.constants.DReplayMemory import MEM_TYPE
+from ai_snake_lab.constants.DLabels import DLabel
 
 
 class AIAgent:
@@ -23,16 +24,13 @@ class AIAgent:
     def __init__(self, epsilon_algo: EpsilonAlgo, seed: int):
         self.epsilon_algo = epsilon_algo
         self.memory = ReplayMemory(seed=seed)
-        self.model = ModelL(seed=seed)
-        # self.model = ModelRNN(seed=seed)
-        self.trainer = AITrainer(self.model)
+        # self._model = ModelL(seed=seed)
+        self._model = ModelRNN(seed=seed)
+        self.trainer = AITrainer(model=self._model)
 
-        if type(self.model) == ModelRNN:
+        if type(self._model) == ModelRNN:
             self.memory.mem_type(MEM_TYPE.RANDOM_GAME)
 
-    def get_model(self):
-        return self.model
-
     def get_move(self, state):
         random_move = self.epsilon_algo.get_move()  # Explore with epsilon
         if random_move != False:
@@ -42,7 +40,7 @@ class AIAgent:
         final_move = [0, 0, 0]
         if type(state) != torch.Tensor:
             state = torch.tensor(state, dtype=torch.float)  # Convert to a tensor
-        prediction = self.model(state)  # Get the prediction
+        prediction = self._model(state)  # Get the prediction
         move = torch.argmax(prediction).item()  # Select the move with the highest value
         final_move[move] = 1  # Set the move
         return final_move  # Return
@@ -50,6 +48,15 @@ class AIAgent:
     def get_optimizer(self):
         return self.trainer.get_optimizer()
 
+    def model_type(self):
+        if type(self._model) == ModelL:
+            return DLabel.MODEL_LINEAR
+        elif type(self._model) == ModelRNN:
+            return DLabel.MODEL_RNN
+
+    def model(self):
+        return self._model
+
     def played_game(self, score):
         self.epsilon_algo.played_game()
 
@@ -57,27 +64,23 @@ class AIAgent:
         # Store the state, action, reward, next_state, and done in memory
         self.memory.append((state, action, reward, next_state, done))
 
-    def set_model(self, model):
-        self.model = model
-
     def set_optimizer(self, optimizer):
         self.trainer.set_optimizer(optimizer)
 
     def train_long_memory(self):
-        # Get the states, actions, rewards, next_states, and dones from the mini_sample
-        memory = self.memory.get_memory()
-        memory_type = self.memory.mem_type()
-
-        if type(self.model) == ModelRNN:
-            for state, action, reward, next_state, done in memory[0]:
-                self.trainer.train_step(state, action, reward, next_state, [done])
-
-        elif memory_type == MEM_TYPE.SHUFFLE:
-            for state, action, reward, next_state, done in memory:
-                self.trainer.train_step(state, action, reward, next_state, [done])
-
-        else:
-            for state, action, reward, next_state, done in memory[0]:
+        # Train on 5 games
+        max_games = 2
+        # Get a random full game
+        while max_games > 0:
+            max_games -= 1
+            game = self.memory.get_random_game()
+            if not game:
+                return  # no games to train on yet
+
+            for count, (state, action, reward, next_state, done) in enumerate(
+                game, start=1
+            ):
+                # print(f"Move #{count}: {action}")
                 self.trainer.train_step(state, action, reward, next_state, [done])
 
     def train_short_memory(self, state, action, reward, next_state, done):

ai_snake_lab/ai/AITrainer.py

@@ -15,23 +15,27 @@ import numpy as np
 import time
 import sys
 
-from ai.models.ModelL import ModelL
-from ai.models.ModelRNN import ModelRNN
+from ai_snake_lab.ai.models.ModelL import ModelL
+from ai_snake_lab.ai.models.ModelRNN import ModelRNN
 
-from constants.DModelL import DModelL
-from constants.DModelLRNN import DModelRNN
+from ai_snake_lab.constants.DModelL import DModelL
+from ai_snake_lab.constants.DModelLRNN import DModelRNN
 
 
 class AITrainer:
 
     def __init__(self, model):
         torch.manual_seed(1970)
-        self.model = model
+
         # The learning rate needs to be adjusted for the model type
         if type(model) == ModelL:
             learning_rate = DModelL.LEARNING_RATE
         elif type(model) == ModelRNN:
             learning_rate = DModelRNN.LEARNING_RATE
+        else:
+            raise ValueError(f"Unknown model type: {type(model)}")
+
+        self.model = model
         self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
         self.criterion = nn.MSELoss()
         self.gamma = 0.9