model.py

# ──────────────────────────────────────────────────────────
#  Project  : Snake_AI
#  File     : model.py
#  Author   : ProGen18
#  Created  : 25-02-2026
#  Modified : 27-03-2026
#  Purpose  : Defines the DQN architecture and the training loop manager
# ──────────────────────────────────────────────────────────


# ============================================================
#                          IMPORTS
# ============================================================

import copy
import os

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau

from config import CONFIG


# ============================================================
#                    NEURAL NETWORK CORE
# ============================================================

# ╔══════════════════════════════════════════════════════════╗
# ║                   ReseauNeurones                         ║
# ╠══════════════════════════════════════════════════════════╣
# ║  Handles:                                                ║
# ║   • Dueling DQN network architecture                     ║
# ║   • Forward pass calculations for state evaluation       ║
# ║   • Saving and loading model weights from disk           ║
# ╚══════════════════════════════════════════════════════════╝


class ReseauNeurones(nn.Module):
    """
    Dueling Q-Network architecture for the Snake AI.

    Attributes:
        features (nn.Sequential): Shared feature extraction layers.
        value (nn.Sequential): Stream estimating the value of the state.
        advantage (nn.Sequential): Stream estimating the advantage of each action.
    """

    def __init__(
        self, input_size: int = CONFIG.input_size, output_size: int = CONFIG.output_size
    ):
        """
        Initialize the neural network layers.

        Args:
            input_size (int): Dimension of the input state vector. Defaults to config input_size.
            output_size (int): Number of possible actions. Defaults to config output_size.
        """
        super().__init__()

        # --- 1. Sizing ---
        c1 = CONFIG.taille_couche_1
        c2 = CONFIG.taille_couche_2
        cv = CONFIG.taille_couche_v
        ca = CONFIG.taille_couche_a

        # --- 2. Shared Features ---
        self.features = nn.Sequential(
            nn.Linear(input_size, c1),
            nn.LayerNorm(c1),
            nn.ReLU(),
            nn.Linear(c1, c2),
            nn.LayerNorm(c2),
            nn.ReLU(),
        )

        # --- 3. Dueling Streams ---
        self.value = nn.Sequential(nn.Linear(c2, cv), nn.ReLU(), nn.Linear(cv, 1))
        self.advantage = nn.Sequential(
            nn.Linear(c2, ca), nn.ReLU(), nn.Linear(ca, output_size)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Calculate Q-values for a given state using Dueling architecture.

        Args:
            x (torch.Tensor): The input state tensor.

        Returns:
            torch.Tensor: The calculated Q-values for all possible actions.
        """
        f = self.features(x)
        v = self.value(f)
        a = self.advantage(f)
        return v + a - a.mean(dim=1, keepdim=True)

    def sauvegarder(
        self,
        nom_fichier: str = "modele.pth",
        nb_parties: int = 0,
        temps_total: float = 0,
        etat_optimiseur: dict | None = None,
        epsilon: float | None = None,
        record: int = 0,
    ) -> None:
        """
        Save the model state and training context to disk.

        Args:
            nom_fichier (str): Name of the file to save. Defaults to 'modele.pth'.
            nb_parties (int): Number of games played. Defaults to 0.
            temps_total (float): Total training time. Defaults to 0.
            etat_optimiseur (dict | None): Optimizer state dictionary. Defaults to None.
            epsilon (float | None): Current exploration rate. Defaults to None.
            record (int): All-time high score. Defaults to 0.
        """
        # --- 1. Directory Setup ---
        dossier = "./model"
        if not os.path.exists(dossier):
            os.makedirs(dossier)

        chemin = os.path.join(dossier, nom_fichier)
        if not chemin.endswith(".pth"):
            chemin += ".pth"

        # --- 2. State Packaging ---
        donnees = {
            "version": 2,
            "etat_modele": self.state_dict(),
            "nb_parties": nb_parties,
            "temps_total": temps_total,
            "etat_optimiseur": etat_optimiseur,
            "epsilon": epsilon,
            "record": record,
        }

        # --- 3. Output ---
        try:
            torch.save(donnees, chemin)
        except Exception as e:
            print(f"Erreur de sauvegarde : {e}")

    def charger(
        self, nom_fichier: str = "modele.pth", device: str = "cpu"
    ) -> tuple | None:
        """
        Load the model state and training context from disk.

        Args:
            nom_fichier (str): Name of the file to load. Defaults to 'modele.pth'.
            device (str): Device to map loaded tensors to. Defaults to 'cpu'.

        Returns:
            tuple | None: Training stats if successful, or None on failure.
        """
        # --- 1. Path Verification ---
        dossier = "./model"
        chemin = os.path.join(dossier, nom_fichier)

        if not os.path.exists(chemin):
            return None

        # --- 2. Checkpoint Loading ---
        try:
            checkpoint = torch.load(chemin, map_location=device)

            if isinstance(checkpoint, dict) and "etat_modele" in checkpoint:
                self.load_state_dict(checkpoint["etat_modele"])
                return (
                    checkpoint.get("nb_parties", 0),
                    checkpoint.get("temps_total", 0),
                    checkpoint.get("etat_optimiseur", None),
                    checkpoint.get("epsilon", None),
                    checkpoint.get("record", 0),
                )
            elif isinstance(checkpoint, dict) and "model_state" in checkpoint:
                # NOTE: Legacy version support
                self.load_state_dict(checkpoint["model_state"])
                return (
                    checkpoint.get("n_games", 0),
                    checkpoint.get("total_time", 0),
                    checkpoint.get("optimizer_state", None),
                    checkpoint.get("epsilon", None),
                    checkpoint.get("record", 0),
                )
            else:
                # Handle bare model dict
                self.load_state_dict(checkpoint)
                return (0, 0, None, None, 0)
        except Exception as e:
            print(f"Erreur chargement : {e}")
            return None


# ============================================================
#                      TRAINING PROCESS
# ============================================================

# ╔══════════════════════════════════════════════════════════╗
# ║                   Entraineur                             ║
# ╠══════════════════════════════════════════════════════════╣
# ║  Handles:                                                ║
# ║   • Learning rate scheduling and optimization            ║
# ║   • Target network soft updates                          ║
# ║   • Q-learning backpropagation and loss calculation      ║
# ╚══════════════════════════════════════════════════════════╝


class Entraineur:
    """
    Manager for the model training process, loss calculation, and optimization.

    Attributes:
        lr (float): Current learning rate.
        gamma (float): Discount factor for future rewards.
        modele (ReseauNeurones): The main policy network.
        tau (float): Rate for target network soft updates.
        device (str): Compute device ('cpu' or 'cuda').
        target_model (ReseauNeurones): Copy of model used to stabilize training.
        optimiseur (optim.Adam): Model optimizer.
        critere (nn.SmoothL1Loss): Loss metric (Huber loss).
        scheduler (ReduceLROnPlateau): Learning rate adjuster.
    """

    def __init__(
        self,
        modele: ReseauNeurones,
        lr: float,
        gamma: float,
        device: str = "cpu",
        tau: float = CONFIG.tau,
    ):
        """
        Initialize the training environment manager.

        Args:
            modele (ReseauNeurones): The neural network to train.
            lr (float): Starting learning rate.
            gamma (float): Discount factor.
            device (str): Device to perform training on. Defaults to 'cpu'.
            tau (float): Target network soft update parameter. Defaults to config target.
        """
        self.lr = lr
        self.gamma = gamma
        self.modele = modele
        self.tau = tau
        self.device = device

        self.target_model = copy.deepcopy(modele).to(device)
        self.target_model.eval()

        self.optimiseur = optim.Adam(modele.parameters(), lr=self.lr)
        self.critere = nn.SmoothL1Loss()

        self.scheduler = ReduceLROnPlateau(
            self.optimiseur,
            mode="max",
            factor=CONFIG.lr_scheduler_factor,
            patience=CONFIG.lr_scheduler_patience,
            min_lr=CONFIG.lr_min,
            verbose=True,
        )

    def mise_a_jour_douce(self) -> None:
        """
        Perform a soft update of the target network parameters using polyak averaging.
        """
        for target_param, local_param in zip(
            self.target_model.parameters(), self.modele.parameters()
        ):
            target_param.data.copy_(
                self.tau * local_param.data + (1.0 - self.tau) * target_param.data
            )

    def etape_d_apprentissage(
        self,
        etat: list | np.ndarray,
        action: list | int,
        recompense: list | float,
        etat_suiv: list | np.ndarray,
        finis: list | bool,
        n_step: int = 1,
    ) -> float:
        """
        Execute a single optimization step using a batch of transitions.

        Args:
            etat (list | np.ndarray): Starting state(s).
            action (list | int): Action(s) taken.
            recompense (list | float): Reward(s) received.
            etat_suiv (list | np.ndarray): Resulting state(s).
            finis (list | bool): Whether the episode ended.
            n_step (int): Number of steps for n-step return. Defaults to 1.

        Returns:
            float: The calculated loss value for the step.
        """
        # --- 1. Validation Setup ---
        etat = torch.tensor(np.array(etat), dtype=torch.float).to(self.device)
        etat_suiv = torch.tensor(np.array(etat_suiv), dtype=torch.float).to(self.device)
        action = torch.tensor(np.array(action), dtype=torch.long).to(self.device)
        recompense = torch.tensor(np.array(recompense), dtype=torch.float).to(
            self.device
        )
        finis = torch.tensor(np.array(finis), dtype=torch.bool).to(self.device)

        if len(etat.shape) == 4:
            etat = etat.view(etat.size(0), -1)
            etat_suiv = etat_suiv.view(etat_suiv.size(0), -1)
        if len(etat.shape) == 1:
            etat = etat.unsqueeze(0)
            etat_suiv = etat_suiv.unsqueeze(0)
            action = action.unsqueeze(0)
            recompense = recompense.unsqueeze(0)
            finis = finis.unsqueeze(0)

        # --- 2. Q-Value Forward Pass ---
        pred = self.modele(etat)

        with torch.no_grad():
            meilleures_actions = self.modele(etat_suiv).argmax(dim=1, keepdim=True)
            next_pred = self.target_model(etat_suiv)
            max_next_q = next_pred.gather(1, meilleures_actions).squeeze(1)

        # NOTE: Bellman equation using target network
        Q_bellman = recompense + self.gamma**n_step * max_next_q * (~finis).float()

        target = pred.clone()
        target.scatter_(1, action.unsqueeze(1), Q_bellman.unsqueeze(1))

        # --- 3. Optimization ---
        self.optimiseur.zero_grad()
        loss = self.critere(pred, target)
        loss.backward()

        torch.nn.utils.clip_grad_norm_(self.modele.parameters(), max_norm=1.0)

        self.optimiseur.step()
        self.mise_a_jour_douce()

        return loss.item()