Résumé Technique : FF-JEPA : Planification à Longue Horizon dans les Modèles de Monde avec des Planificateurs Latents

Énoncé du Problème

Les avancées récentes dans les modèles de monde visuels, particulièrement les architectures de prédiction d'encastrement conjoint (JEPA) comme LeWorldModel (LeWM), ont démontré une capacité à prédire les observations futures et à soutenir le contrôle basé sur des modèles via des planificateurs par échantillonnage comme la méthode de l'entropie croissante (CEM). Cependant, les auteurs identifient trois limitations critiques empêchant le déploiement de ces modèles de monde « plats » dans des tâches réelles à longue horizon :

1. Prohibitions Computationnelles : La planification sur de longs horizons nécessite des rollouts répétés et des séquences d'actions de plus en plus grandes, rendant l'optimisation par CEM computationnellement intraitable.
2. Accumulation d'Erreurs : Les erreurs de prédiction autoregressive s'accumulent sur de nombreuses étapes, provoquant la divergence du planificateur pour des trajectoires complexes.
3. Dépendance au But : Les méthodes existantes nécessitent généralement une image ou un état de but explicite ($o_g$) pour piloter la planification. Dans de nombreux scénarios réels, une spécification de but aussi concrète est indisponible ou peu pratique.

Bien que les méthodes hiérarchiques et les approches génératives aient tenté d'aborder ces problèmes, elles opèrent souvent encore dans l'espace des pixels ou dépendent d'images/trajectoires de buts spécifiées de l'extérieur, limitant leur applicabilité dans des environaux ouverts.

Méthodologie : Forward-Forward JEPA (FF-JEPA)

Le papier propose FF-JEPA, un cadre hiérarchique qui étend les modèles de monde de type JEPA en introduisant un planificateur latent ($G$) opérant aux côtés d'un modèle de propagation standard conditionné par l'action ($P$). Cette approche décompose la planification à long horizon en une séquence de problèmes d'optimisation à court terme tractables sans nécessiter d'images de but explicites.

Architecture Centrale

Le système se compose de deux composants principaux entraînés dans l'espace latent défini par un encodeur de modèle de monde pré-entraîné et gelé ($E$) :

1. Le Modèle de Monde ($P$) : Un modèle de dynamique de propagation standard conditionné par l'action (basé sur LeWM) qui prédit l'état latent suivant étant donné un historique d'états et une action : $\hat{z}_{t+1} = P(z_{t-W_P+1:t}, a_t)$.
2. Le Planificateur Latent ($G$) : Un modèle de dynamique de propagation sans action qui prédit le prochain sous-but latent ($\hat{z}_{sg}$) basé uniquement sur un historique de sous-buts latents précédents. Il opère à une résolution temporelle plus élevée (tous les $H$ pas de l'environnement).
   • Planificateur Déterministe ($G_{Det}$) : Une architecture transformer (similaire au prédicteur du modèle de monde mais sans conditionnement par l'action) entraînée pour minimiser l'erreur quadratique moyenne entre les sous-buts prédits et réels.
   • Planificateur de Diffusion ($G_{DM}$) : Un backbone DiT (Diffusion Transformer) entraîné avec un objectif de mise en correspondance de score de débruitage pour prédire une séquence de futurs sous-buts.

Inférence de Planification

Le processus d'inférence est hiérarchique :

1. Prédiction de Sous-But : Le planificateur latent $G$ prédit le prochain état de sous-but $\hat{z}_{sg, m+1}$ étant donné l'historique actuel des sous-buts.
2. Optimisation d'Action : Le modèle de monde $P$ utilise la méthode CEM pour trouver une séquence d'actions ($a_{1:H}$) qui conduit l'état prédit vers le sous-but prédit $\hat{z}_{sg, m+1}$ dans un horizon court $H$.
3. Itération : Ce cycle se répète, enchaînant efficacement des plans à court horizon pour atteindre des objectifs à long horizon.

Crucialement, cette méthode élimine le besoin d'une image de but cible lors de l'inférence ; le planificateur génère la trajectoire vers un objectif implicitement défini.

Contributions Clés

• Planification à Long Horizon sans But : Le cadre permet de planifier dans des modèles de monde sans images de but explicites en utilisant un planificateur latent pour générer des sous-buts intermédiaires.
• Décomposition Hiérarchique : Il traite l'« effondrement à long horizon » des modèles de monde plats en décomposant des trajectoires complexes en sous-problèmes plus courts et solubles, atténuant l'accumulation d'erreurs et les coûts computationnels.
• Entraînement Non Étiqueté : Le planificateur latent peut être entraîné sur des démonstrations réussies non étiquetées (sous-échantillonnées à partir de l'espace latent du modèle de monde), réduisant la dépendance aux données étiquetées par l'action.
• Cadre Unifié : Il réinterprète le modèle de monde comme un module de dynamique inverse opérant sur des trajectoires latentes imaginées, unifiant la modélisation prédictive et le contrôle au sein d'un seul espace latent cohérent.

Évaluation Empirique

Les auteurs ont évalué FF-JEPA sur la tâche PushT à travers trois configurations expérimentales : Horizon court (25 étapes), Long horizon (75 étapes) et Initialisation Aléatoire (300 étapes avec des positions de départ aléatoires).

Résultats

• Performance à Long Horizon : Le modèle de base plat LeWM s'est effondré sur les tâches à long horizon (taux de succès de 3,52 % à $t=75$) et a totalement échoué sous initialisation aléatoire (0,00 %). En revanche, FF-JEPA (Diffusion) a atteint 91,80 % de succès à $t=75$ et 82,42 % sous initialisation aléatoire.
• Avantage du "Sans But" : Contrairement aux bases (LeWM, DINO-WM) qui nécessitent une image de but pour l'inférence, FF-JEPA n'en a pas besoin. Malgré cela, il a surpassé les bases hiérarchiques DINO-WM tant en horizon court (96,09 % contre 89,0 %) qu'en long horizon (91,80 % contre 61,0 %).
• Efficacité des Données : Les études d'ablation ont montré que l'entraînement du planificateur de diffusion sur seulement 200 démonstrations de haute qualité (une réduction de 40x des données) produisait toujours des performances compétitives (76,17 % contre 82,42 % avec les données complètes), suggérant une grande efficacité des données.
• Surcoût Computationnel : Le planificateur déterministe a ajouté un surcoût de temps d'inférence négligeable (2,1 ms) par rapport à CEM (926 ms), tandis que le planificateur de diffusion a ajouté ~242 ms. Les deux sont nettement plus rapides que les rollouts répétés requis pour la planification plate à long horizon.

Analyse d'Échec

Les cas d'échec se sont généralement produits lorsque l'agent s'éloignait trop de la position du sous-but prédit ou lorsque le planificateur latent générait un état de sous-but de faible qualité (par exemple, en omettant l'agent).

Signification et Revendications

Le papier positionne FF-JEPA comme une direction prometteuse pour l'apprentissage robotique scalable et sans but explicite. En démontrant que les modèles de monde conditionnés par l'action couplés à des planificateurs latents entraînés sur des données non étiquetées peuvent surmonter les limites de la planification plate, les auteurs suggèrent que cette approche permet aux agents d'opérer dans des environnements ouverts où les spécifications de buts explicites sont indisponibles. Le travail souligne que la décomposition de la planification en sous-buts à court terme est une stratégie viable pour atténuer les problèmes de calcul et d'accumulation d'erreurs inhérents aux approches actuelles de modèles de monde.