model.py

"""
Deep Q-Network (DQN) Model and Trainer for Snake AI

This module contains the neural network architecture and training logic
for the Snake AI agent using Deep Q-Learning.
"""

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os
from typing import Union, Tuple
from pathlib import Path


class Linear_QNet(nn.Module):
    """
    Simple Deep Q-Network with one hidden layer
    
    Architecture: Input -> Hidden (ReLU) -> Output
    - Input: 11 features about game state
    - Hidden: 256 neurons with ReLU activation  
    - Output: 3 Q-values (one for each action)
    """
    
    def __init__(self, input_size: int, hidden_size: int, output_size: int):
        super().__init__()
        self.input_layer = nn.Linear(input_size, hidden_size)
        self.output_layer = nn.Linear(hidden_size, output_size)
        
        # Initialize weights for better training stability
        self._initialize_weights()
    
    def _initialize_weights(self):
        """Initialize network weights using Xavier initialization"""
        nn.init.xavier_uniform_(self.input_layer.weight)
        nn.init.xavier_uniform_(self.output_layer.weight)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through the network"""
        x = F.relu(self.input_layer(x))
        x = self.output_layer(x)  # No activation on output (raw Q-values)
        return x
    
    def save(self, filename: str = 'best_model.pth') -> None:
        """Save the trained model to disk"""
        save_dir = Path('./saved_models')
        save_dir.mkdir(exist_ok=True)  # Create directory if it doesn't exist
        
        filepath = save_dir / filename
        torch.save(self.state_dict(), filepath)
        print(f"Model saved to {filepath}")
    
    def load(self, filename: str = 'best_model.pth') -> bool:
        """Load a previously trained model from disk"""
        filepath = Path('./saved_models') / filename
        
        if not filepath.exists():
            print(f"Model file {filepath} not found")
            return False
        
        try:
            self.load_state_dict(torch.load(filepath))
            print(f"Model loaded from {filepath}")
            return True
        except Exception as e:
            print(f"Error loading model: {e}")
            return False


class QTrainer:
    """
    Trainer for the Deep Q-Network using the Bellman equation
    
    Implements the core DQN learning algorithm:
    Q(s,a) = r + γ * max(Q(s',a')) for non-terminal states
    Q(s,a) = r for terminal states
    """
    
    def __init__(self, model: Linear_QNet, learning_rate: float, discount_factor: float):
        self.learning_rate = learning_rate
        self.discount_factor = discount_factor  # How much we value future rewards
        self.model = model
        self.optimizer = optim.Adam(model.parameters(), lr=learning_rate)
        self.loss_function = nn.MSELoss()
    
    def train_step(self, state: Union[torch.Tensor, tuple], action: Union[torch.Tensor, tuple], 
                   reward: Union[torch.Tensor, tuple], next_state: Union[torch.Tensor, tuple], 
                   game_over: Union[bool, tuple]) -> float:
        """
        Perform one training step using the DQN algorithm
        
        Returns the loss value for monitoring training progress
        """
        # Convert everything to tensors
        state = torch.tensor(state, dtype=torch.float)
        next_state = torch.tensor(next_state, dtype=torch.float)
        action = torch.tensor(action, dtype=torch.long)
        reward = torch.tensor(reward, dtype=torch.float)
        
        # Handle both single samples and batches
        if len(state.shape) == 1:
            # Single sample - add batch dimension
            state = torch.unsqueeze(state, 0)
            next_state = torch.unsqueeze(next_state, 0)
            action = torch.unsqueeze(action, 0)
            reward = torch.unsqueeze(reward, 0)
            game_over = (game_over, )  # Convert to tuple for consistent handling
        
        # Get current Q-values for all actions
        current_q_values = self.model(state)
        
        # Calculate target Q-values using Bellman equation
        target_q_values = current_q_values.clone()
        
        for batch_idx in range(len(game_over)):
            new_q_value = reward[batch_idx]
            
            if not game_over[batch_idx]:
                # If game continues, add discounted future reward
                future_q_values = self.model(next_state[batch_idx])
                new_q_value = reward[batch_idx] + self.discount_factor * torch.max(future_q_values)
            
            # Update Q-value for the action that was taken
            action_idx = torch.argmax(action[batch_idx]).item()
            target_q_values[batch_idx][action_idx] = new_q_value
        
        # Train the network
        self.optimizer.zero_grad()
        loss = self.loss_function(target_q_values, current_q_values)
        loss.backward()
        self.optimizer.step()
        
        return loss.item()