Résumé technique : Modèles de monde sensorimoteurs (SMWM)

Énoncé du problème

Le papier traite du défi de l'apprentissage de modèles de monde latents à partir d'observations de haute dimension (par exemple, des pixels) sans signaux de récompense ou d'étiquettes de tâches. Bien que les architectures de prédiction d'encapsulation conjointe (JEPA) offrent un cadre prometteur en prédisant directement les états futurs dans un espace d'encapsulation (évitant la reconstruction de pixels), elles souffrent d'un effondrement de la représentation (representation collapse). Lorsqu'un encodeur et un modèle de dynamique directe sont entraînés conjointement pour minimiser l'erreur de prédiction, l'encodeur peut apprendre à mapper toutes les observations vers une encapsulation constante. Cette solution triviale minimise la perte de prédiction mais rend le modèle inutile pour le contrôle ou la planification en aval.

Les solutions existantes pour prévenir l'effondrement reposent souvent sur des mécanismes complexes tels que :

• L'utilisation d'encodeurs pré-entraînés gelés (ex. : DINO-WM).
• L'ajout de régulariseurs distributionnels (ex. : faire correspondre les encapsulations à une gaussienne isotrope via SIGReg).
• L'utilisation de réseaux cibles par moyenne mobile exponentielle (EMA) avec arrêt des gradients (stop-gradients).
• L'emploi de fonctions de perte multi-termes impliquant des contraintes de variance-covariance.

Les auteurs postulent que ces approches ne traitent peut-être pas pleinement la question fondamentale de ce qu'une représentation de modèle de monde utile devrait conserver. S'appuyant sur le concept de « perception pour l'action » et de perception énactive, le papier soutient que les représentations devraient être façonnées par leur pertinence vis-à-vis des actions, préservant les distinctions qui comptent pour les dynamiques contrôlables tout en éliminant les distracteurs incontrôlables.

Méthodologie : Modèle de monde sensorimoteur (SMWM)

Les auteurs proposent le SMWM, un modèle de monde latent entraîné de bout en bout à partir de trajectoires hors ligne (offline) $(o_t, a_t, o_{t+1})$. L'innovation centrale est l'utilisation de la régularisation par dynamique inverse comme unique mécanisme pour prévenir l'effondrement de la représentation.

Architecture

Le modèle se compose de trois composants entraînés conjointement :

1. Encodeur ($f_\theta$) : Mappe les observations $o_t$ vers une encapsulation latente compacte $z_t$.
2. Modèle de dynamique directe ($g_\phi$) : Prédit l'état latent suivant $\hat{z}_{t+1}$ étant donné l'état actuel $z_t$ et l'action $a_t$.
3. Modèle de dynamique inverse ($h_\psi$) : Prédit l'action $\hat{a}_t$ effectuée entre deux encapsulations consécutives $z_t$ et $z_{t+1}$.

Objectif d'entraînement

Le modèle est optimisé en minimisant une fonction de perte conjointe :
$$\mathcal{L} = \mathcal{L}_{fwd} + \lambda \mathcal{L}_{inv}$$

• Perte directe ($\mathcal{L}_{fwd}$) : Erreur quadratique moyenne (MSE) entre l'encapsulation suivante prédite $\hat{z}_{t+1} = g_\phi(z_t, a_t)$ et l'encapsulation réelle de l'état suivant $z_{t+1} = f_\theta(o_{t+1})$.
• Perte inverse ($\mathcal{L}_{inv}$) : MSE entre l'action prédite $\hat{a}_t = h_\psi(z_t, z_{t+1})$ et l'action réelle $a_t$.
• Poids de régularisation ($\lambda$) : Un unique hyperparamètre contrôlant la force du signal de dynamique inverse.

Mécanisme d'action

La perte inverse agit comme un mécanisme anti-effondrement. Si l'encodeur s'effondre vers une encapsulation constante, le modèle inverse reçoit des entrées identiques pour toutes les transitions et ne peut prédire qu'une action constante (la moyenne de la distribution d'actions). Pour réduire la perte inverse en dessous de ce risque de prédicteur constant, l'encodeur doit préserver l'information concernant l'action $a_t$ contenue dans la transition $(o_t, o_{t+1})$. Cela force la représentation à conserver les informations pertinentes pour l'action (degrés de liberté contrôlables) tout en filtrant naturellement les distracteurs incontrôlables (bruit exogène) qui n'influencent pas la récupération de l'action.

Contrairement aux régulariseurs distributionnels qui prescrivent la géométrie de l'espace d'encapsulation (par exemple, en forçant une forme gaussienne), la dynamique inverse ancre la représentation à une quantité ancrée dans la tâche (l'action) sans imposer de priors géométriques spécifiques.

Contributions clés

1. Architecture SMWM : Un modèle de monde sensorimoteur simple, entraînable de bout en bout, qui utilise la dynamique inverse comme unique mécanisme anti-effondrement, ne nécessitant qu'un seul hyperparamètre supplémentaire ($\lambda$).
2. Propriétés de la représentation : Preuves empiriques que les représentations apprises ne sont pas seulement non-effondrées, mais suivent spécifiquement les degrés de liberté contrôlables, préservent la géométrie spatiale (tant cartésienne que périodique) et filtrent les distracteurs incontrôlables.
3. Performance en aval : Démonstration d'une performance de planification compétitive à travers des tâches de contrôle en 2D et 3D, montrant que les avantages structurels de la représentation se traduisent par une planification efficace orientée vers des buts.

Résultats expérimentaux

Bancs d'essai contrôlés (Dot World & Sprite World)

• Récupération de la dimension intrinsèque : Dans un environnement « dot world », le modèle a réussi à récupérer la véritable dimension intrinsèque 2D de l'espace d'état à partir des pixels seuls, même lorsqu'il était intégré dans un espace latent de 64 dimensions.
• Filtrage des distracteurs : Lorsqu'on introduisait des objets « distracteurs » incontrôlables (se déplaçant de manière aléatoire), l'encodeur les ignorait. La dimension latente effective correspondait aux degrés de liberté contrôlables, et non à la dimension totale de l'observation.
• Perception dépendante du contrôle : Dans une expérience de « sprite world », le modèle a appris à préserver les variables de pose (position, orientation) uniquement lorsqu'elles étaient contrôlables. Lorsque l'orientation n'était pas contrôlée, la représentation la faisait la moyenne, résultant en une reconstruction floue de l'orientation.

Benchmarks de planification

Les auteurs ont évalué SMSM sur quatre tâches de contrôle hors ligne : TwoRoom (navigation 2D), Reacher (contrôle continu), Push-T (manipulation 2D) et OGBench-Cube (manipulation 3D).

• Baselines : Comparaison avec SIGReg (une baseline JEPA forte utilisant des régulariseurs gausiens), Forward-only (sans perte inverse) et l'échantillonnage d'actions Random.
• Performance :
  • SMWM a égalé la performance de SIGReg sur les trois tâches 2D.
  • SMSM a surpassé SIGReg sur la tâche 3D OGBench-Cube (84 % de succès contre 59 %).
  • Les deux méthodes régularisées ont nettement surpassé l'ablation Forward-only, confirmant la nécessité d'un mécanisme anti-effondrement.
• Géométrie latente : Des expériences de sondage ont révélé que les encapsulations de SMWM forment des sous-espaces de faible dimension reflétant la géométrie physique de l'état (par exemple, des structures cylindriques pour les angles de joints périodiques dans Reacher, des cartes cartésiennes pour la navigation).

Signification et affirmations

Le papier affirme que SMWM fournit une instanciation simple et fondée sur des principes de la « perception pour l'action » dans l'apprentissage de modèles de monde latents. En forçant les états latents à préserver l'information suffisante pour récupérer les actions de l'agent, le modèle biaise naturellement la représentation vers les degrés de liberté contrôlables de l'environnement.

Les points clés incluent :

• Simplicité : La dynamique inverse seule suffit à stabiliser l'entraînement end-to-end de JEPA sans reconstruction, sans signaux de récompense, sans encodeurs gelés, ni priors distributionnels complexes.
• Biais structurel : La méthode induit des représentations qui sont interprétables et alignées avec la structure causale de l'environnement (actions et leurs effets), plutôt que de simples corrélations prédictives.
• Utilité pratique : Les représentations apprises permettent une planification latente efficace, particulièrement dans des environnements 3D complexes où les autres baselines régularisées peinent.

Les auteurs reconnaissent des limites, telles que l'hypothèse que les actions sont récupérables à partir d'observations consécutives et l'incapacité actuelle à capturer des quantités non identifiables à partir de trames uniques (comme la vitesse) sans historiques d'observation. Cependant, ils positionnent SMWM comme une approche robuste et sans récompense pour apprendre des modèles de monde actionnables à partir de données hors ligne.