init: tetris neural network model with q learning

README.md

@@ -0,0 +1,87 @@
+# Tetris-Neural-Network-Q-Learning
+
+
+## Overview
+**PyTorch** implementation of a simplified Tetris-playing AI using **Q-Learning**.  
+The Tetris board is just 4×4, with the agent deciding in which of the 4 columns to drop the next piece. The agent’s neural network receives a **16-dimensional** board representation (flattened 4×4) and outputs **4** Q-values, one for each possible move. Through repeated training (via self-play and the Q-Learning algorithm), the agent learns to fill the board without making illegal moves—eventually achieving a perfect score.
+
+<img src="images/game.png" />
+
+## Project Structure
+
+```plaintext
+
+├── model.py            # Contains the TetrisAI class and TetrisNet model (PyTorch)
+├── train.py            # Main training script
+├── evaluate.py         # Script to load the model checkpoint and interactively run the game
+├── tetris.py           # Defines the GameState and game logic
+├── representation.py   # Defines how the game state is turned into a 1D list of ints
+└── checkpoints         # Directory where model checkpoints (.pth) are saved/loaded
+```
+
+## Model Architecture
+- **Input Layer (16 units):** Flattened 4x4 board state, where each cell is `0` (empty) or `1` (occupied).
+- **Hidden Layers:** Dense layers (64 → 64 → 32) with ReLU activations.
+- **Output Layer (4 units):** Linear activation, representing the estimated Q-value for each move (column 1–4).
+
+## Training
+1. **Game Environment:** A 4x4 Tetris-like grid where each move places a block in one of the four columns.
+2. **Reward Function:** 
+   - **Immediate Reward:** Increase in the number of occupied squares, minus 
+   - **Penalty:** A scaled standard deviation of the “column depth” to encourage balanced play.
+3. **Q-Learning Loop:**
+   - For each move, the model is passed the current game state and returns predicted Q-values.
+   - An action (move) is chosen based on either:
+     - **Exploitation:** Highest Q-value prediction (greedy choice).
+     - **Exploration:** Random move to discover new states.
+   - The agent observes the new state and reward, and stores this experience (state, action, reward, next_state) to update the model.
+
+## Reward Function
+
+The reward function for each action is based on two parts:
+
+1. **Board Occupancy**
+   - The reward starts with the number of occupied squares on the board (i.e., how many cells are filled).
+
+2. **Penalty for Unbalanced Columns**
+   - Next, the standard deviation of each column's unoccupied cell count is calculated.
+   - A higher standard deviation means one column may be much taller or shorter than others, which is undesirable in Tetris.
+   - By *subtracting* this standard deviation from the occupancy-based reward, the agent is penalized for building unevenly and is encouraged to keep the board as level as possible.
+
+In other words:
+
+\[
+\text{Reward} = \text{OccupiedSquares} - \alpha \times \text{StdDev}(\text{ColumnDepths})
+\]
+
+Where \( \alpha \) is a weighting factor (in this case effectively 1, or any scalar you choose) that determines the penalty's intensity. This keeps the board balanced and helps the agent learn a more efficient Tetris strategy.
+
+## Installation & Usage
+1. Clone this repo or download the source code.
+2. Install Python (3.8+ recommended).
+3. Install dependencies:
+
+    ```bash
+    pip install torch numpy
+    ```
+   - You may also need other libraries like pandas or statistics depending on your environment.
+
+1. Training:
+
+   - Adjust the hyperparameters (learning rate, exploration rate, etc.) in ```train.py``` if desired.
+   - Run:
+
+```bash
+python train.py
+```
+
+   - This script will generate a ```checkpointX.pth``` file in checkpoints/ upon completion (or periodically during training).
+
+1. Evaluation:
+
+   - Ensure you have a valid checkpoint saved, for example ```checkpoint14.pth.```
+   - Run:
+    ```bash
+    python evaluate.py
+    ```
+   - The script will load the checkpoint, instantiate the ```TetrisAI```, and then interactively show how the AI plays Tetris. You can step through the game move by move in the console.

evaluate.py

@@ -0,0 +1,41 @@
+import model
+import tetris
+import representation
+from pathlib import Path
+
+script_dir = Path(__file__).parent.resolve()
+checkpoints_dir = script_dir / "checkpoints"
+checkpoints_dir.mkdir(parents=True, exist_ok=True)
+checkpoint_filename = "checkpoint14.pth"
+save_path = checkpoints_dir / checkpoint_filename
+
+# If you need it as a standard Python string:
+save_path_str = str(save_path)
+
+tai = model.TetrisAI(save_path)
+
+while True:
+    gs = tetris.GameState()
+    while True:
+
+        print("Board:")
+        print(str(gs))
+
+        # get move
+        predictions:list[float] = tai.predict(representation.BoardState(gs))
+        shift:int = predictions.index(max(predictions))
+        print("Move: " + str(shift))
+        input("Enter to execute the move it selected: ")
+
+        # make move
+        gs.drop(shift)
+
+        # if game over
+        if gs.over():
+            print(str(gs))
+            print("Game is over!")
+            print("Final score: " + str(gs.score()))
+            print("Going to next game...")
+            gs = tetris.GameState()
+            gs.randomize()
+

model.py

@@ -0,0 +1,102 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import numpy as np
+
+class Experience:
+    def __init__(self):
+        self.state: list[int] = None
+        self.action: int = None
+        self.reward: float = None
+        self.next_state: list[int] = None
+        self.done: bool = False
+
+class TetrisNet(nn.Module):
+    """
+    The PyTorch neural network equivalent to your Keras model:
+    Input: 16-dimensional board
+    Hidden layers: 64 -> 64 -> 32, ReLU activation
+    Output: 4-dimensional, linear
+    """
+    def __init__(self):
+        super(TetrisNet, self).__init__()
+        self.layer1 = nn.Linear(16, 64)
+        self.layer2 = nn.Linear(64, 64)
+        self.layer3 = nn.Linear(64, 32)
+        self.output = nn.Linear(32, 4)
+        self.relu = nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.relu(self.layer1(x))
+        x = self.relu(self.layer2(x))
+        x = self.relu(self.layer3(x))
+        x = self.output(x)
+        return x
+
+class TetrisAI:
+    """
+    PyTorch implementation of the TetrisAI class.
+    - Loads a saved model if save_file_path is provided.
+    - Otherwise, constructs a fresh model.
+    - Has methods to save, predict, and train the model.
+    """
+
+    def __init__(self, save_file_path: str = None):
+        # Create the model
+        self.model = TetrisNet()
+
+        # Define the optimizer and loss function
+        self.optimizer = optim.Adam(self.model.parameters(), lr=0.003)
+        self.criterion = nn.MSELoss()
+
+        # Load from file if path is provided
+        if save_file_path is not None:
+            checkpoint = torch.load(save_file_path, map_location=torch.device('cpu'))
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+            self.model.eval()
+
+    def save(self, path: str) -> None:
+        """
+        Saves the PyTorch model and optimizer state to a file.
+        """
+        torch.save({
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict()
+        }, path)
+
+    def predict(self, board: list[int]) -> list[float]:
+        """
+        Performs a forward pass to predict the Q-values for each possible move.
+        Returns these Q-values as a list of floats.
+        """
+        # Convert board to a float tensor with shape [1, 16]
+        x = torch.tensor([board], dtype=torch.float32)
+
+        # Put model in evaluation mode and disable gradient tracking
+        self.model.eval()
+        with torch.no_grad():
+            prediction = self.model(x)
+
+        # Convert the single batch output (shape [1, 4]) to a Python list of floats
+        return prediction[0].tolist()
+
+    def train(self, board: list[int], qvalues: list[float]) -> None:
+        """
+        Trains the model on one step using the given board as input and qvalues as the desired output.
+        """
+        # Put model in training mode
+        self.model.train()
+
+        # Convert data to tensors
+        x = torch.tensor([board], dtype=torch.float32)
+        y = torch.tensor([qvalues], dtype=torch.float32)
+
+        # Zero the parameter gradients
+        self.optimizer.zero_grad()
+
+        # Forward + Backward + Optimize
+        predictions = self.model(x)
+        loss = self.criterion(predictions, y)
+        loss.backward()
+        self.optimizer.step()

play.py

@@ -0,0 +1,20 @@
+import tetris
+
+while True:
+    gs = tetris.GameState()
+    
+    while True:
+        print("Board:")
+        print(str(gs))
+
+        i:str = input("How many shifts? > ")
+        shifts:int = int(i)
+        reward = gs.drop(shifts)
+        print("REWARD: " + str(reward))
+
+        # if game over
+        if gs.over():
+            print("Game over!")
+            print("Score: " + str(gs.score()))
+            input("Enter to go to next game.")
+            gs = tetris.GameState()

representation.py

@@ -0,0 +1,9 @@
+import tetris
+
+def BoardState(gs:tetris.GameState) -> list[int]:
+    """Represents the board as a state of flattened integers."""
+    ToReturn:list[int] = []
+    for row in gs.board:
+        for col in row:
+            ToReturn.append(int(col))
+    return ToReturn

tetris.py

@@ -0,0 +1,121 @@
+import statistics
+import random
+
+class InvalidDropException(Exception):
+    def __init__(self, message):
+        self.message = message
+        super().__init__(self.message)
+
+class GameState:
+    def __init__(self):
+        self.board: list[list[bool]] = [
+            [False, False, False, False],
+            [False, False, False, False],
+            [False, False, False, False],
+            [False, False, False, False],
+        ]  # 4 rows of 4 columns, 4x4
+
+    def __str__(self):
+        ToReturn: str = ""
+        ToReturn = " ┌────┐" + "\n"
+        onRow: int = 0
+        for row in self.board:
+            # add the row number in
+            ToReturn = ToReturn + str(onRow) + "│"
+
+            # print every square
+            for column in row:
+                if column:
+                    ToReturn = ToReturn + "█"
+                else:
+                    ToReturn = ToReturn + " "
+            ToReturn = ToReturn + "│\n"
+            onRow = onRow + 1
+        ToReturn = ToReturn + " └────┘"
+        ToReturn = ToReturn + "\n" + "  0123"
+        return ToReturn
+
+    def column_depths(self) -> list[int]:
+        """Calculates how 'deep' the available space on each column goes, from the top down."""
+
+        # record the depth of every column
+        column_depths: list[int] = [0, 0, 0, 0]
+        column_collisions: list[bool] = [
+            False,
+            False,
+            False,
+            False,
+        ]
+
+        # In this sense, "depth" is the number of squares that are clear, to be clear
+        for ri in range(0, len(self.board)):  # for every row
+            for ci in range(
+                0, len(self.board[0])
+            ):  # for every column (use first row to know how many columns there are)
+                if (
+                    column_collisions[ci] == False and self.board[ri][ci] == False
+                ):  # if column X has not been recorded yet and the column in this row is not occupied, increment the depth
+                    column_depths[ci] = column_depths[ci] + 1
+                else:  # we hit a floor!
+                    column_collisions[ci] = True
+
+        return column_depths
+
+    def over(self) -> bool:
+        """Determines the game is over (if all cols in top row are occupied)."""
+        return self.board[0] == [1, 1, 1, 1]
+
+    def drop(self, column: int) -> float:
+        """Drops a single block into the column, returns the reward of doing so."""
+        if column < 0 or column > 3:
+            raise InvalidDropException(
+                "Invalid move! Column to drop in must be 0, 1, 2, or 3."
+            )
+
+        reward_before: float = self.score_plus()
+        cds: list[int] = self.column_depths()
+        if cds[column] == 0:
+            raise InvalidDropException(
+                "Unable to drop on column " + str(column) + ", it is already full!"
+            )
+        self.board[cds[column] - 1][column] = True
+        reward_after: float = self.score_plus()
+        return reward_after - reward_before
+
+    def score(self) -> int:
+        ToReturn: int = 0
+        for row in self.board:
+            for col in row:
+                if col:
+                    ToReturn = ToReturn + 1
+        return ToReturn
+
+    def score_plus(self) -> float:
+        # start at score
+        ToReturn: float = float(self.score())
+
+        # penalize for standard deviation
+        stdev: float = statistics.pstdev(self.column_depths())
+        ToReturn = ToReturn - (stdev * 2)
+
+        return ToReturn
+
+    def randomize(self) -> float:
+        """Sets the board to a random setup."""
+
+        # first, clear all values
+        for ri in range(0, len(self.board)):
+            for ci in range(0, len(self.board[0])):
+                self.board[ri][ci] = False
+
+        # drop a random number in each column
+        for ci in range(0, 4):
+            random_drops: int = random.randint(0, 4)
+            for _ in range(0, random_drops):
+                self.drop(ci)
+
+        # if all 16 are filled up, delete one
+        if self.score() == 16:
+            self.board[0][random.randint(0, 3)] = (
+                False  # turn off a random square in the top row
+            )

train.py

@@ -0,0 +1,141 @@
+import model
+import tetris
+import sys
+import representation
+import random
+from pathlib import Path
+
+script_dir = Path(__file__).parent.resolve()
+checkpoints_dir = script_dir / "checkpoints"
+checkpoints_dir.mkdir(exist_ok=True) 
+log_file_path = checkpoints_dir / "log.txt"
+
+# if you want to start from a checkpoint, fill this in with the path to the .pth file. If wanting to start from a new NN, leave blank!
+model_save_path = r"" 
+
+# training settings
+gamma:float = 0.5
+epsilon:float = 0.2
+
+# training config
+batch_size:int = 100 # the number of experiences that will be collected and trained on
+save_model_every_experiences:int = 5000
+################
+
+# construct/load model
+tmodel:model.TetrisAI = None
+if model_save_path != None and model_save_path != "":
+    print("Loading model checkpoint at '" + model_save_path + "'...")
+    tmodel = model.TetrisAI(model_save_path)
+    print("Model loaded!")
+else:
+    print("Constructing new model...")
+    tmodel = model.TetrisAI()  
+
+# variables to track
+experiences_trained:int = 0 # the number of experiences the model has been trained on
+model_last_saved_at_experiences_trained:int = 0 # the last number of experiences that the model was trained on
+on_checkpoint:int = 0
+
+def log(path:str, content:str) -> None:
+    if path != None and path != "":
+        f = open(path, "a")
+        f.write(content + "\n")
+        f.close()
+        
+# training loop
+while True:
+
+    # collect X number of experiences
+    gs:tetris.GameState = tetris.GameState()
+    experiences:list[model.Experience] = []
+    for ei in range(0, batch_size):
+
+        # print!
+        sys.stdout.write("\r" + "Collecting experience " + str(ei+1) + " / " + str(batch_size) + "... ")
+        sys.stdout.flush()
+
+        # get board representation
+        state_board:list[int] = representation.BoardState(gs)
+
+        # select move to play
+        move:int
+        if random.random() < epsilon: # if by chance we should select a random move
+            move = random.randint(0, 3) # choose move at random
+        else:
+            predictions:list[float] = tmodel.predict(state_board) # predict Q-Values
+            move = predictions.index(max(predictions)) # select the move (index) with the highest Q-Value
+
+        # play the move
+        IllegalMovePlayed:bool = False
+        MoveReward:float
+        try:
+            MoveReward = gs.drop(move)
+        except tetris.InvalidDropException as ex: # the model (or at random) tried to play an illegal move
+            IllegalMovePlayed = True
+            MoveReward = -3.0 # small penalty for illegal moves
+        except Exception as ex:
+            print("Unhandled exception in move execution: " + str(ex))
+            input("Press enter key to continue, if you want to.")
+        
+        # store this experience
+        exp:model.Experience = model.Experience()
+        exp.state = state_board
+        exp.action = move
+        exp.reward = MoveReward
+        exp.next_state = representation.BoardState(gs) # the state we find ourselves in now.
+        exp.done = gs.over() or IllegalMovePlayed # it is over if the game is completed OR an illegal move was played
+        experiences.append(exp)
+
+        # if game is over or they played an illegal move, reset the game!
+        if gs.over() or IllegalMovePlayed:
+            gs = tetris.GameState()
+            
+    print()
+
+    # print avg rewards
+    rewards:float = 0.0
+    for exp in experiences:
+        rewards = rewards + exp.reward
+    status:str = "Average reward over those " + str(len(experiences)) + " experiences on model w/ " + str(experiences_trained) + " trained experiences: " + str(round(rewards / len(experiences), 2))
+    log(log_file_path, status)
+    print(status)
+    
+    # train!
+    for ei in range(0, len(experiences)):
+        exp = experiences[ei]
+
+        # print training number
+        sys.stdout.write("\r" + "Training on experience " + str(ei+1) + " / " + str(len(experiences)) + "... ")
+        sys.stdout.flush()
+
+        # determine new target based on the game ending or not (maybe we should factor in future rewards, maybe we shouldnt)
+        new_target:float
+        if exp.done:
+            new_target = exp.reward
+        else:
+            max_q_of_next_state:float = max(tmodel.predict(exp.next_state))
+            new_target = exp.reward + (gamma * max_q_of_next_state) # blend immediate vs. future rewards
+
+        # ask the model to predict again for this experiences state
+        qvalues:list[float] = tmodel.predict(exp.state)
+
+        # plug in the new target where it belongs
+        qvalues[exp.action] = new_target
+
+        # now train on the updated qvalues (with 1 changed)
+        tmodel.train(exp.state, qvalues)
+        experiences_trained = experiences_trained + 1
+        
+    print("Training complete!")
+
+    # save model!
+    if (experiences_trained - model_last_saved_at_experiences_trained) >= save_model_every_experiences:
+        print("Time to save model!")
+        path = checkpoints_dir / f"checkpoint{on_checkpoint}.pth"
+        
+        tmodel.save(path)
+        print("Checkpoint # " + str(on_checkpoint) + " saved to " + str(path) + "!")
+        on_checkpoint = on_checkpoint + 1
+        model_last_saved_at_experiences_trained = experiences_trained
+        print("Model saved to " + str(path) + "!")