Fondamentaux de l'apprentissage par renforcement (RL)

L'apprentissage par renforcement repose sur l'interaction entre deux entités principales : l'Agent et l'Environnement. Cette dynamique s'articule autour de trois concepts clés :

• Espace d'état (S) : L'ensemble des situations possibles dans lesquelles l'environnement peut se trouver.
• Espace d'action (A) : L'ensemble des décisions ou mouvements que l'agent peut effectuer.
• Récompense (R) : Le signal scalaire renvoyé par l'environnement suite à une action, indiquant la qualité du passage d'un état à un autre.

Le cycle de fonctionnement est le suivant : à l'instant $t$, l'agent observe l'état $S_t$ et reçoit une récompense $R_t$. Il choisit une action $A_t$, ce qui modifie l'état de l'environnement vers $S_{t+1}$ et génère une nouvelle récompense $R_{t+1}$. L'objectif ultime de l'agent est de maximiser la somme pondérée des récompenses futures, appelée Valeur ($V$).

L'équation de valeur simplifiée peut s'écrire :

Dans le contexte du RLHF (Reinforcement Learning from Human Feedback), l'objectif est d'ajuster un modèle de langage pour qu'il génère des réponses conformes aux préférences humaines. Voici la correspondance des concepts :

• Agent : Le modèle de langage (LLM).
• Action ($A_t$) : La génération d'un jeton (token) spécifique à un instant donné. L'espace d'action est le vocabulaire du modèle.
• État ($S_t$) : Le contexte actuel, comprenant l'instruction initiale (prompt) et tous les jetons générés précédemment.
• Environnement : Le processus de génération de texte.

L'architecture à quatre modèles du PPO

La mise en œuvre de PPO pour l'alignement des modèles de langage mobilise généralement quatre réseaux distincts :

1. Actor Model (Acteur) : C'est le modèle cible que nous entraînons. Initialisé à partir d'un modèle SFT (Supervised Fine-Tuning), son but est de produire la réponse optimale.
2. Reference Model (Référence) : Une copie figée du modèle SFT. Il sert de garde-fou pour éviter que l'Acteur ne s'éloigne trop de ses capacités linguistiques d'origine (prévention de l'effondrement du modèle).
3. Critic Model (Critique) : Un modèle entraîné pour estimer la valeur attendue $V_t$. Il juge la qualité potentielle de la séquence complète à partir de l'état actuel.
4. Reward Model (Récompense) : Un modèle figé qui simule le jugement humain en attribuant un score scalaire à la séquence générée.

Mécanismes de calcul de la récompense et de la perte

La perte de l'Acteur ne repose pas uniquement sur le score du Reward Model. Pour stabiliser l'apprentissage, on intègre une pénalité basée sur la divergence KL entre l'Acteur et le modèle de Référence.

def generer_recompenses(self, jetons_entree, log_probs_actuels, log_probs_ref, score_brut, masque_action):
    """
    Calcule la récompense combinant le score du Reward Model et la pénalité KL.
    """
    # Calcul de la pénalité KL pour maintenir la cohérence avec le modèle initial
    # divergence_kl = log(P_acteur) - log(P_ref)
    diff_kl = log_probs_actuels - log_probs_ref
    recompenses_base = -self.kl_coefficient * diff_kl
    
    # On identifie la fin de la séquence générée
    index_depart = jetons_entree.shape[1] - 1
    indices_fin = index_depart + masque_action[:, index_depart:].sum(1) + 1
    
    # Limitation de l'amplitude du score de récompense
    score_borne = torch.clamp(score_brut, -self.limite_recompense, self.limite_recompense)
    
    batch_size = log_probs_actuels.shape[0]
    for i in range(batch_size):
        # On ajoute le score du Reward Model uniquement sur le dernier token de la réponse
        recompenses_base[i, index_depart:indices_fin[i]][-1] += score_borne[i]
        
    return recompenses_base

Estimation de l'avantage (GAE)

Pour mettre à jour les poids, PPO utilise la notion d'Avantage ($Adv$). L'avantage mesure si une action a été meilleure ou pire que l'estimation moyenne faite par le Critique. Le GAE (Generalized Advantage Estimation) permet de lisser ces estimations.

def estimer_avantages_et_retours(self, valeurs, recompenses, debut_reponse):
    """
    Calcule l'avantage GAE et les retours cibles pour le Critique.
    """
    avantage_cumule = 0
    liste_avantages = []
    taille_seq = recompenses.size()[-1]
    
    # Parcours inversé pour la programmation dynamique du GAE
    for t in reversed(range(debut_reponse, taille_seq)):
        valeur_suivante = valeurs[:, t + 1] if t < taille_seq - 1 else 0.0
        
        # Erreur temporelle (TD-error)
        delta_td = recompenses[:, t] + self.gamma * valeur_suivante - valeurs[:, t]
        
        # Accumulation de l'avantage avec le facteur lambda
        avantage_cumule = delta_td + self.gamma * self.lam * avantage_cumule
        liste_avantages.append(avantage_cumule)
        
    avantages = torch.stack(liste_avantages[::-1], dim=1)
    retours_reels = avantages + valeurs[:, debut_reponse:]
    
    return avantages.detach(), retours_reels

Optimisation de la fonction de perte Actor-Critic

Le cœur de PPO réside dans sa fonction de perte "clippée". Lors d'un cycle d'entraînement, on réutilise le même lot de données (batch) pendant plusieurs époques. Pour éviter que le modèle ne change trop radicalement en une seule mise à jour, on contraint le ratio entre la nouvelle politique et l'ancienne.

La perte de l'Acteur est définie par :

# Double calcul pour l'écrêtage
surr1 = -avantages * ratio
surr2 = -avantages * torch.clamp(ratio, 1.0 - self.clip_range, 1.0 + self.clip_range)

# On prend le maximum (car on cherche à minimiser l'opposé de l'objectif)
perte_finale = torch.sum(torch.max(surr1, surr2) * masque) / masque.sum()
return perte_finale

Parallèlement, le modèle Critique est mis à jour pour minimiser l'erreur quadratique entre ses prédictions et les retours réels calculés via GAE, en appliquant également un mécanisme d'écrêtage pour stabiliser l'apprentissage de la fonction de valeur.