Résumé technique : Modèles de monde latents à longueur variable pour la planification à long horizon

1. Énoncé du problème

Les modèles de monde sont apparus comme un paradigme pour construire des agents intelligents en apprenant des modèles prédictifs de la dynamique de l'environnement. Bien que les modèles de monde latents récents (par exemple, les architectures de type JEPA) offrent des alternatives efficaces à la prédiction dans l'espace des pixels en opérant dans des espaces de représentation compacts, ils font face à des limitations significatives en matière de planification à long horizon.

Les principaux problèmes identifiés sont :

• Erreurs cumulées : Les modèles existants reposent généralement sur une prédiction à un pas ($t \to t+1$). Pour les tâches à long horizon, ces modèles doivent être déployés de manière récursive, ce qui provoque une accumulation rapide des erreurs de prédiction, entraînant une dérive temporelle et une dégradation de la qualité de la trajectoire.
• Décalage entre entraînement et planification : La plupart des modèles sont entraînés avec un objectif de prédiction à pas fixe. Cela ne fournit une supervision qu'à une échelle temporelle locale, échouant à optimiser explicitement la cohérence à long terme. Par conséquent, le signal d'entraînement ne s'aligne pas sur la nature des problèmes de planification à long horizon.
• Rigidité architecturale : Les modèles de monde latents actuels incorporent souvent les actions via des mécanismes de conditionnement dédiés (par exemple, des couches FiLM ou l'attention croisée) qui fusionnent les actions dans les états cachés. Cette conception lie étroitement le modèle à une structure d'action fixe (généralement une action par étape), ce qui rend difficile la distinction des contributions de multiples actions ou la gestion naturelle de séquences d'actions de longueurs variables.

2. Méthodologie : Modèles de monde latents à longueur variable (VLWM)

Les auteurs proposent VLWM, un cadre conçu pour apprendre la dynamique de l'environnement latent à travers plusieurs horizons de prédiction au sein d'un seul modèle.

2.1 Supervision à longueur variable

Au lieu de restreindre la supervision à l'étape suivante immédiate, VLWM entraîne un prédicteur unique $g_\phi$ pour prévoir l'état latent à un décalage futur arbitraire $k \ge 1$. L'objectif d'entraînement est :
$$\mathcal{L}_{VLWM} = \mathbb{E}_{t, k \sim p(k)} \left[ \| g_\phi(s_t, a_{t:t+k-1}) - \bar{g}(\bar{f}_{\bar{\theta}}(o_{t+k})) \|_2^2 \right]$$
où $s_t$ est l'état latent actuel, $a_{t:t+k-1}$ est un segment d'action de longueur variable $k$, et la cible est l'état latent à l'instant $t+k$. Cela unifie la dynamique à court terme et la cohérence sémantique à long terme sous un objectif unique.

2.2 Formulation de l'action en tant que jeton (Action-as-Token)

Pour supporter nativement les séquences d'actions de longueur variable, VLWM s'écarte du conditionnement interne aux couches. Au lieu de cela, il traite les plongements d'actions comme des jetons (tokens) de premier rang du transformer.

• Séquence d'entrée : Le modèle reçoit une séquence $X$ contenant un historique de jetons d'état suivis d'un nombre variable de jetons d'action ($e_t, \dots, e_{t+k-1}$).
• Mécanisme : Les jetons d'état et d'action sont distingués par des plongements de type additifs. Le transformer est causal, et le latent prédit $\hat{s}_{t+k}$ est lu directement à partir de la sortie du dernier jeton de la séquence.
• Avantage : Cela élimine le besoin d'interfaces de conditionnement d'action prédéfinies, permettant au modèle de se conditionner sur un nombre arbitraire d'actions passées sans changement architectural.

2.3 Stratégie d'entraînement par curriculum

L'entraînement direct avec de grands horizons de prédiction dès le début est instable en raison de gradients bruités et d'échecs d'optimisation. VLWM emploie une stratégie d'apprentissage par curriculum :

• L'entraînement est divisé en étapes. Dans l'étape $j$, l'horizon $k$ est échantillonné uniformément dans $\{1, 2, \dots, j\}$.
• Le modèle commence avec $k=1$ et élargit progressivement la plage des horizons échantillonnés jusqu'à un maximum $K_{max}$.
• Cela garantit que la sémantique à long terme est construite sur des dynamiques à court horizon bien apprises, stabilisant ainsi l'optimisation.

2.4 Planification avec des sauts de longueur variable

Lors de l'inférence, l'horizon de planification $H$ est partitionné en blocs contigus $H = k_1 + k_2 + \dots + k_M$. Le planificateur évalue le déploiement comme une séquence de sauts latents de longueur variable plutôt que comme des prédictions récursives à un pas. Le papier propose trois stratégies de planification :

1. P1 (Fixe mais adaptatif à la tâche) : Utilise une taille de bloc constante $k^*$ sélectionnée par tâche (petite pour un contrôle fin, grande pour des objectifs à long horizon).
2. P2 (Échantillonnage d'horizon aléatoire) : Échantillonne des tailles de blocs uniformément dans $\{1, \dots, K_{max}\}$ à chaque frontière.
3. P3 (Du long vers le court) : Commence par de grands blocs pour couvrir la structure à long terme et réduit progressivement la taille des blocs vers la fin pour affiner les dynamiques fines (une recherche du grossier vers le fin).

3. Contributions clés

1. Cadre unifié : Introduction de VLWM, un cadre qui apprend la dynamique latente à travers plusieurs horizons et supporte naturellement les séquences d'actions de longueur variable via une formulation "action-en-tant-que-jeton".
2. Entraînement et planification stables : Proposition d'une stratégie d'apprentissage par curriculum évolutive pour un entraînement à longueur variable stable et d'une suite d'algorithmes de planification qui exploitent de manière adaptative différents horizons de prédiction au sein d'un seul modèle.
3. Validation empirique : Démonstration que VLWM surpasse systématiquement les modèles de monde à horizon fixe (spécifiquement LeWM) sur des benchmarks de planification à long horizon, avec des analyses montrant une meilleure prédictibilité à long terme et une meilleure contrôlabilité en aval.

4. Résultats expérimentaux

Les auteurs ont évalué VLWM sur trois environnements de planification conditionnés par un objectif : PushT (manipulation 2D), OGBench-Cube (manipulation robotique 3D) et TwoRoom (navigation 2D).

• Gains de performance : VLWM a obtenu une amélioration moyenne de 13 % par rapport à l'état de l'art LeWM sur les jeux de données. Les gains ont été les plus significatifs dans les régimes à long horizon (par exemple, sur TwoRoom avec un horizon de 100 étapes, le taux de réussite est passé de 36 % avec LeWM à 68 % avec VLWM).
• Atténuation des erreurs : Les visualisations ont montré que tandis que les déploiements récursifs de LeWM dérivaient et échouaient à franchir les frontières environnementales (par exemple, les portes), les sauts par blocs de VLWM restaient proches des trajectoires réelles.
• Qualité de la représentation : Les évaluations par sondage linéaire sur les encodeurs gelés ont révélé que VLWM produisait des latents avec une corrélation de Pearson plus élevée pour les propriétés physiques (par exemple, l'angle d'un bloc) par rapport à LeWM, suggérant une meilleure capture des changements d'état sémantiques lents.
• Nécessité du curriculum : Les études d'ablation ont confirmé que la stratégie d'entraînement par curriculum est essentielle pour un entraînement à longueur variable stable, particulièrement sur les tâches nécessitant une planification étendue.

5. Signification et affirmations

Le papier affirme que VLWM fournit un paradigme simple mais efficace pour améliorer la prédiction et la planification à long horizon dans les modèles de monde latents.

• Au-delà de l'ingénierie : Les auteurs soutiennent que la difficulté de la planification à long horizon n'est pas simplement un problème d'ingénierie, mais une conséquence des objectifs d'entraînement. En passant d'une supervision à pas fixe à une supervision à longueur variable, VLWM aligne le signal d'entraînement avec les exigences du raisonnement à long horizon.
• Interface unifiée : Le cadre subsume un éventail de stratégies à horizon fixe ou hiérarchique au sein d'un seul modèle entraîné, éliminant le besoin de politiques distinctes de haut niveau et de bas niveau ou de modules auxiliaires.
• Direction future : Les résultats suggèrent que la supervision à longueur variable est une étape fondamentale vers des algorithmes de planification plus robustes, offrant une interface flexible où les planificateurs peuvent adapter la granularité de la prédiction aux besoins spécifiques de la tâche.