Résumé Technique : CAST (Credit Assignment from Solver Teachers)

Énoncé du Problème

L'entraînement des modèles de langage à grande échelle (LLM) pour agir en tant qu'agents généralistes dans des environaux à horizon long, tels que les jeux classiques (Sokoban, Minesweeper, Rush Hour), se heurte à un goulot d'étranglement critique : la rareté des signaux de récompense. Si l'apprentissage par renforcement avec récompenses vérifiables (RLVR) a réussi pour les tâches à tour unique, il repose sur des récompenses de résultat final (victoire/défaite) qui ne fournissent aucune information sur les décisions intermédiaires spécifiques ayant conduit au succès ou à l'échec. Cela crée un problème fondamental d'attribution de crédit où chaque jeton d'une trajectoire reçoit le même retour grossier, entravant la capacité du modèle à apprendre la planification à long terme, l'exploration orientée vers un objectif et la récupération d'erreurs. Les tentatives existantes pour densifier ces signaux via une recherche coûteuse, des modèles de récompense de processus appris ou des comparaisons entre trajectoires souffrent souvent de compromis entre coût computationnel, qualité de la supervision et fiabilité du signal.

Méthodologie : CAST

Les auteurs proposent CAST (Credit Assignment from Solver Teachers), un cadre qui exploite des solveurs de jeux spécifiques au domaine pour fournir une supervision fine, au niveau du tour, aux agents LLM.

Mécanisme Central

1. Le Solveur comme Enseignant au Niveau du Tour : Au lieu d'utiliser un solveur uniquement pour générer des trajectoires expertes (ce qui limite l'apprentissage aux états visités par l'expert), CAST utilise le solveur pour évaluer les actions on-policy du LLM en temps réel.
2. Construction de l'Avantage du Solveur :
   • Le solveur calcule une valeur de coût à venir (cost-to-go), $N(s)$, représentant le nombre minimal d'étapes nécessaires pour gagner à partir de l'état $s$.
   • Pour une action $a_t$ entreprise par le LLM transitant de $s_t$ vers $s_{t+1}$, la méthode calcule le changement du coût à venir.
   • Un avantage de solveur décalé est défini comme :
     $$\tilde{A}^{\pi_{Solver}}(s_t, a_t) = N(s_t) - \mathbb{E}[N(s_{t+1})]$$
   • Cette valeur est positive si l'action réduit la distance vers l'objectif (progrès), nulle si elle est neutre, et négative si elle augmente la distance ou mène à une impasse.
3. Façonnage du Signal (Signal Shaping) : Pour éviter l'instabilité de l'entraînement causée par des valeurs à queue lourde (ex: fortes pénalités pour les impasses) et les différences d'échelle entre les jeux, CAST applique deux transformations :
   • Compression Asinh : $g(x) = \text{asinh}(x)$, qui préserve les petites valeurs de manière linéaire tout en compressant les valeurs aberrantes extrêmes de manière logarithmique.
   • Normalisation RMS au niveau du Batch : Rescale le signal par sa magnitude racine carrienne moyenne (RMS) au sein d'un batch pour assurer des échelles cohérentes, sans soustraire la moyenne (pour préserver la signification sémantique de zéro comme étant "aucun progrès").
4. Intégration avec RLVR : L'avantage du solveur façonné est ajouté à l'avantage standard de trajectoire de résultat (ex: issu de GRPO) pour former un signal d'avantage combiné :
   $$\hat{A}_{i,t} = \hat{A}^{outcome}_i + \alpha \cdot h(\tilde{A}^{\pi_{Solver}}_{i,t})$$
   Ce signal combiné est utilisé dans l'objectif de substitut tronqué (clipped surrogate objective) pour mettre à jour la politique.

Insight Théorique : Distillation On-Policy Sans Logits

Le papier établit une équivalence théorique entre la maximisation de l'avantage du solveur et la distillation on-policy (OPD) à partir d'un enseignant "soft-optimal".

• Sous l'hypothologie qu'un solveur fort agit comme une politique soft-optimale où les probabilités d'action sont proportionnelles à $\exp(Q/\tau)$, l'avantage du solveur est mathématiquement proportionnel à la log-probabilité de l'action sous la politique du solveur.
• Contribution Clé : Cela implique que l'avantage scalaire du solveur encode intrinsèquement la distribution de préférence de l'enseignant. Par conséquent, CAST réalise la distillation sans nécessiter les logits de l'enseignant (la distribution d'action complète), en s'appuyant uniquement sur des valeurs scalaires. Cela permet au LLM étudiant d'apprendre de la fonction de valeur du solveur efficacement.

Contributions Clés

1. Crédit Dérivé du Solveur : L'introduction d'une méthode pour convertir les changements de valeur d'état des solveurs de jeux en avantages fins, au niveau du tour, pour l'entraînement des LLM.
2. Distillation Sans Logits : Une preuve théorique que la maximisation des avantages du solveur est équivalente à la distillation on-policy, permettant un entraînement efficace utilisant uniquement des valeurs scalaires plutôt que des distributions d'actions complètes.
3. Façonnage de Signal Stable : L'utilisation de la transformation asinh et de la normalisation RMS par batch pour gérer les valeurs extrêmes et les variations d'échelle entre domaines, assurant une intégration stable dans le RLVR.
4. Praticité : Démonstration que les solveurs exacts ajoutent un surcoût négligeable et que les réseaux de valeur appris (ex: DQN) peuvent servir d'approximations efficaces lorsqu'aucun solveur exact n'est disponible.

Résultats Expérimentaux

La méthode a été évaluée sur Sokoban, Minesweeper et Rush Hour, avec un transfert zero-shot sur ALFWorld et WebShop.

• Performance In-Domain : CAST a surpassé tous les baselines entraînés (GRPO, GSPO, DAPO, GiGPO) et les méthodes de prompting sans entraînement (incluant des modèles fermés puissants comme Claude-Opus-4.6) dans les trois jeux, tant dans les réglages in-domain que dans les réglages de difficulté non vue.
  • En moyenne, CAST a atteint un taux de succès de 62,1 % (Avg@4) sur les tâches in-domain et 28,4 % sur les difficultés non vues, surpassant significativement le backbone DAPO (44,7 % ID, 18,7 % Unseen).
• Efficacité d'Échantillonnage : CAST a atteint le pic de performance de validation de DAPO en 1,7× à 2,0× moins d'étapes d'entraînement à travers les trois jeux.
• Généralisation Zero-Shot : Sans ajustement fin supplémentaire, les agents CAST ont été transférés vers ALFWorld et WebShop, atteignant la performance moyenne zero-shot la plus élevée (30,3 % global) parmi toutes les méthodes, dépassant de 5,6 points le second meilleur baseline.
• Études d'Ablation :
  • Le poids de l'avantage du solveur ($\alpha$) est critique ; $\alpha=0,1$ a fourni le meilleur équilibre entre guidage de processus et objectifs de résultat.
  • Tant la compression asinh que la normalisation RMS étaient essentielles ; la suppression de l'une ou l'autre dégradait les performances, la suppression de l'asinh causant la chute la plus importante due à l'instabilité du gradient face aux valeurs extrêmes.
• Solveurs Appris : Remplacer le solveur exact par un réseau de valeur DQN appris conserve la majeure partie des bénéfices de performance, démontrant l'applicabilité de la méthode au-delà des jeux possédant des solveurs exacts.

Signification et Revendications

Le papier affirme que CAST répond au défi fondamental de l'attribution de crédit dans le RLVR à long horizon en fournissant un signal fin et scalable dérivé de solveurs de domaine. En prouvant l'équivalence avec la distillation on-policy sans logits, la méthode offre une voie pratique pour aligner les LLM avec des fonctions de valeur expertes sans la charge computationnelle de la génération de distributions d'actions complètes. Les résultats suggèrent qu'une évaluation d'état fiable — que ce soit via des solveurs exacts ou des fonctions de valeur apprises — offre une route robuste passant de résultats épars à un feedback dense et actionnable, permettant aux LLM de maîtriser des tâches de prise de décision complexes à plusieurs étapes plus efficacement que les approches basées uniquement sur le résultat.