Résumé Technique : Modèles de Monde Équivariants au Flux (Flow Equivariant World Models)

1. Énoncé du Problème

Les agents incarnés opèrent dans des environnements dynamiques où ils ne possèdent qu'un champ de vision de première personne restreint. Cette observabilité partielle crée un défi majeur pour les modèles de monde existants : alors que l'agent observe une fraction de l'environnement, les régions non observées continuent d'évoluer. Les méthodes de pointe actuelles, incluant les transformeurs de diffusion latente à grande échelle et les modèles récurrents, échouent souvent à maintenir la cohérence sur de longs horizons dans ces contextes.

Sans une structure de mémoire respectant les symétries sous-jacentes du monde, ces modèles souffrent de :

• Perte d'information : À mesure que l'agent se déplace, les informations passées pertinentes sortent de la fenêtre de contexte (par exemple, les fenêtres glissantes des transformeurs) et sont jetées.
• Hallucination : Lorsqu'un agent se retourne vers une région précédemment observée, le modèle échoue souvent à se rappeler l'état de l'objet, générant à la place de nouveaux objets hallucinés ou faisant disparaître les objets existants.
• Absence de priors structurels : Les modèles existants traitent les observations comme des séquences statiques ou des données dépendantes de la vue, ignorant le fait que tant le mouvement propre de l'agent que le mouvement externe des objets obéissent à des symétries lisses paramétrées par le temps (flux).

2. Méthodologie : Modèles de Monde Équivariants au Flux (FloWM)

Les auteurs proposent FloWM, un cadre qui intègre l'équivariance au flux au cœur de la mémoire du modèle de monde. L'hypothèse centrale est que les représentations prédictives deviennent plus puissantes lorsqu'elles sont organisées en alignement avec la structure temporelle et dynamique du monde.

Concept Central : Équivariance au Flux

Le cadre étend l'équivariance de groupe statique aux transformations de séquences paramétrées par le temps, ou « flux ».

• Définition : Un modèle est équivariant au flux si, lorsque la séquence d'entrée subit un flux (par exemple, un décalage dû au temps ou au mouvement), la séquence de sortie se transforme selon ce même flux.
• Équivariance Globale : Dans les contextes d'observation partielle, le modèle doit maintenir une mémoire latente de l'état entier du monde (observé et non observé). La représentation latente doit se transformer de manière équivariante par rapport à l'état global du monde, et non seulement par rapport à l'observation actuelle.
• Référentiel Co-mobile : Pour y parvenir, la mémoire récurrente est transportée par le même flux que l'entrée avant d'incorporer de nouvelles observations. Cela garantit que les caractéristiques restent alignées spatialement avec les objets en mouvement plutôt que d'être laissées derrière dans un référentiel fixe.

Composantes Architecturales

Le cadre est instancié par une relation de récurrence généralisée équivariante au flux :
$$h_{t+1}(\nu) = \psi_1(\nu) \cdot U_\theta [h_t(\nu); E_\theta [f_t; h_t] (\nu)]$$
Où :

• $h_t(\nu)$ représente l'état caché organisé en « canaux de vélocité » ($\nu$), correspondant à différents champs de vecteurs de mouvement.
• $E_\theta$ est un encodeur mappant les observations vers l'espace latent.
• $U_\theta$ est l'opération de mise à jour.
• $\psi_1(\nu)$ représente l'action du flux sur un pas de temps.

Équivariance au Mouvement Propre :
Le modèle traite le mouvement propre de l'agent comme un flux relatif. En transformant l'état caché selon l'action de l'agent ($T_{a_t}$), la mémoire se déplace efficacement dans le référentiel co-mobile de l'agent. Cela permet au modèle de construire une « carte » structurée de l'environnement où les nouvelles observations sont intégrées dans les positions relatives correctes en fonction des actions de l'agent.

Deux Instanciations

1. FloWM Récurrent Simple (2D) : Conçu pour les environnements 2D (ex: MNIST World). Il utilise des opérations convolutionnelles sur les états cachés et les entrées spatiales. Il emploie une fenêtre fixe pour lire/écrire dans l'état caché tandis que le reste de l'état circule autour du champ de vision de l'agent.
2. FloWM basé sur un Transformer (3D) : Conçu pour des environnements 3D complexes (ex: Block World). Il utilise un encodeur et un décodeur Vision Transformer (ViT). La mémoire latente est un ensemble de plongements de tokens organisés spatialement agissant comme une carte abstraite descendante.
   • Note sur l'équivariance : Bien que l'équivariance exacte de l'encodeur soit difficile pour les projections 3D vers 2D sans une dé-projection de profondeur coûteuse, les auteurs démontrent que la transformation de flux récurrente ($T^{-1}_{a_t} \psi_1(\nu)$) encourage l'encodeur à apprendre l'équivariance pendant l'entraînement, atteignant une haute fidélité sans contraintes géométriques explicites.

3. Principales Contributions

• Formalisation de l'Équivariance au Flux : Le papier formalise comment les symétries paramétrées par le temps (flux) peuvent être utilisées pour construire une mémoire structurée et persistante pour la modélisation de mondes dynamiques partiellement observés.
• Architecture de Mémoire Structurée : Introduction d'une mémoire récurrente avec des « canaux de vélocité » qui se déplace et se transforme de manière équivariante avec le mouvement propre et le mouvement externe des objets, maintenant l'information sur les régions hors champ alignée au fil du temps.
• Nouveaux Benchmarks : Les auteurs introduisent de nouveaux benchmarks pour la modélisation de mondes partiellement observés :
  • MNIST World 2D : Chiffres se déplaçant à des vélocités aléatoires avec observation partielle.
  • Block World 3D Dynamique : Blocs colorés rebondissant dans une pièce 3D avec un agent naviguant via des actions discrètes.
• Validation Empirique : Démonstration que FloWM surpasse significativement les modèles de diffusion (DFoT, DFoT-SSM) et les modèles récurrents (RSSM) de pointe en termes de prédiction sur de longs horizons et de généralisation de longueur.

4. Résultats Expérimentaux

Les auteurs ont évalué FloWM par rapport à des bases telles que le History-Guided Diffusion Forcing (DFoT), DFoT-SSM et le RSSM de DreamerV3.

Métriques de Performance

• Cohérence sur Long Horizon : Sur le MNIST World 2D, FloWM a maintenu des prédictions cohérentes pendant 150 pas de temps au-delà de l'horizon d'entraînement (20 pas de temps), là où les bases divergeaient ou hallucinaient. Sur le Block World 3D, FloWM est resté cohérent pendant 210 images, tandis que les bases perdaient la trace des objets ou hallucinaient de nouveaux éléments.
• Réduction de l'Erreur : FloWM a obtenu une erreur quadratique moyenne (MSE) nettement plus faible et des scores PSIM/SSIM plus élevés par rapport à toutes les bases. Par exemple, dans le Block World 3D, l'MSE de FloWM était d'environ 0,0006 (70 images) contre 0,0117 pour DFoT.
• Efficacité d'Apprentissage : Les modèles combinant l'Équivariance au Mouvement Propre (SME) et les Canaux de Vélocité (VC) ont nécessité des ordres de grandeur moins d'étapes d'entraînement pour converger par rapport aux versions ablatives.
• Planification Aval : Dans une tâche de planification « Trouver le bloc rouge », FloTM a réussi à naviguer vers la cible. Les bases ont souvent halluciné l'emplacement de la cible, entraînant une mauvaise performance de planification.
• Analyse par Sondage (Probe) : Un réseau de sondage entraîné pour décoder les positions des blocs à partir de l'espace latent a montré une précision de 96 % pour FloWM contre <1 % pour DFoT, confirmant que FloWM apprend une représentation structurée et équivariante des trajectoires d'objets.

Études d'Ablation

Le retrait des Canaux de Vélocité (VC) ou de l'Équivariance au Mouvement Propre (SME) a provoqué une dégradation significative des performances, rejoignant ou tombant même en dessous des modèles de diffusion de base, soulignant que ces deux composants sont critiques pour gérer l'observabilité partielle et la dynamique.

5. Signification et Revendications

Le papier affirme que les représentations prédictives sont fondamentalement plus puissantes lorsqu'elles sont organisées en accord avec la structure temporelle et dynamique du monde.

• Fondements Théoriques : Ce travail fournit un cadre théorique (équivariance au flux) qui dicte la nécessité d'une mémoire structurée pour les modèles de monde dans des environnements dynamiques et partiellement observés.
• Résolution de l'Observabilité Partielle : Contrairement aux méthodes précédentes qui reposent sur l'extension des fenêtres de contexte ou des banques de mémoire statiques, FloWM gère intrinsèquement l'évolution des régions non observées en maintenant une mémoire qui se déplace avec l'état du monde.
• Généralisation : Le cadre permet de généraliser à des séquences nettement plus longues que celles vues durant l'entraînement, une capacité critique pour les agents incarnés du monde réel.
• Portée Modeste : Les auteurs reconnaissent que les évaluations actuelles se déroulent dans des contextes contrôlés avec des actions géométriques de corps rigides. Ils positionnent ce travail comme une fondation pour de futures extensions vers des jeux de données en monde ouvert, des espaces d'actions continus et des interactions physiques plus complexes.

En résumé, FloWM démontre qu'en encodant explicitement les symétries du mouvement dans l'architecture de la mémoire, les modèles de monde peuvent atteindre des prédictions stables, précises et sans hallucinations sur de longs horizons dans des environnements dynamiques partiellement observés.