Résumé Technique : Modèle du Monde Prismatique (PRISM-WM)

1. Énoncé du Problème

L'apprentissage par renforcement basé sur un modèle (MBRL) et la planification dans les domaines robotiques sont fondamentalement limités par la nature hybride des dynamiques physiques. Les systèmes réels, en particulier dans la locomotion bipède et la manipulation, impliquent un mouvement continu ponctué d'événements discrets (par exemple, contacts, impacts, adhérence vs glissement). Ces interactions créent des dynamiques non lisses et hautement non linéaires, avec des discontinuités abruptes dans le champ vectoriel sous-jacent.

Les modèles de monde latents conventionnels emploient généralement des réseaux de neurones monolithiques qui imposent une continuité globale. Bien que théoriquement capables d'approximation universelle, ces modèles apprennent en pratique une représentation moyenne qui lisse excessivement les modes dynamiques distincts. Pour les planificateurs d'optimisation de trajectoire, ce lissage excessif induit des erreurs cumulatives lors des anticipations à long horizon, en particulier aux frontières physiques. Cela se traduit par :

• Un flou spectral à travers les frontières de contact, causant un désaccord significatif des gradients.
• Une dérive cumulative dans les espaces de configuration de haute dimension, rendant la planification à long terme peu fiable.
• Un effondrement de la politique dans les environnements riches en contacts, où les modèles standards échouent à distinguer les régimes stables des régimes instables.

2. Méthodologie : Le Modèle du Monde Prismatique (PRISM-WM)

Pour remédier à ces limitations, les auteurs introduisent PRISM-WM, une architecture structurée conçue pour décomposer les dynamiques hybrides complexes en primitives composables. Le modèle est construit sur une base de Mélange d'Experts (MoE) conscient du contexte avec un objectif novateur d'orthogonalisation latente.

Architecture Principale

• Gating Conscient du Contexte : Un réseau de gating ($G$) reçoit un vecteur de contexte $c_t$ (comprenant l'état $z_t$, l'action $a_t$, et optionnellement l'embedding de tâche $e_{task}$) et produit des poids normalisés par softmax $w_t$. Ce mécanisme identifie implicitement le régime physique actif (par exemple, vol vs appui) sans supervision explicite.
• Experts Spécialisés : L'état est acheminé vers un ensemble de $K$ réseaux experts $\{E_k\}$. Chaque expert se spécialise dans l'apprentissage du flux continu au sein d'un régime dynamique local spécifique, évitant ainsi les interférences avec des modes conflictuels.
• Prédiction Résiduelle : Le prochain état latent est prédit comme une combinaison pondérée des sorties des experts ajoutée à l'état actuel :
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  où $v_k$ représente les dynamiques résiduelles prédites par le $k$-ième expert.

Innovation Clé : Orthogonalisation Latente

Pour prévenir l'effondrement des modes (où les experts apprennent des caractéristiques redondantes) et assurer la diversité, PRISM-WM introduit une couche d'orthogonalisation de Gram-Schmidt différentiable.

• Après que chaque expert a produit un vecteur de la couche pénultième $u_k$, le modèle applique un processus de Gram-Schmidt stabilisé pour générer des vecteurs de base orthogonaux $v_k$.
• Cela force les représentations des experts à se situer sur une variété de Stiefel, garantissant qu'elles englobent un sous-espace de rang maximal et non redondant pour les dynamiques résiduelles.
• Bien que les modes de contact physiques soient non ordonnés, cette orthogonalisation assure que les experts capturent des composantes distinctes de la variété résiduelle, permettant au réseau de gating de reconstruire des modes physiques non ordonnés via une somme pondérée.

Intégration avec la Planification

PRISM-WM est intégré dans le cadre de planification TD-MPC. En factorisant les dynamiques complexes en experts non redondants, le modèle réduit le résidu du modèle en un pas de temps de $\epsilon$ à environ $\epsilon/K\alpha$. Cette décomposition structurelle atténue la dérive cumulative, permettant des déroulements (rollouts) à long horizon de haute fidélité et une propagation de gradient stable pour l'apprentissage de la politique.

3. Contributions Clés

1. Décomposition Architecturale : Introduction de PRISM-WM, qui décompose les dynamiques hybrides en utilisant une architecture MoE consciente du contexte couplée à une orthogonalisation latente, s'éloignant des approximations globales monolithiques.
2. Objectif d'Orthogonalisation : Un objectif d'entraînement novateur utilisant l'orthogonalisation de Gram-Schmidt différentiable pour imposer la diversité des experts et prévenir l'effondrement des modes, alignant structurellement l'espace latent avec les transitions physiques discrètes.
3. Validation Théorique et Empirique : Démonstration que cette décomposition réduit les erreurs cumulatives dans les espaces de haute dimension et améliore la géométrie du paysage de perte tant pour la planification (TD-MPC) que pour l'apprentissage direct de la politique (PWM).
4. Capacité de Transfert Zero-Shot : Preuve que la base orthogonale permet au modèle de généraliser à travers des simulateurs avec des physiques de contact différentes (par exemple, de IsaacGym à contact souple à MuJoCo à contact rigide) en isolant des modes dynamiques spécifiques (par exemple, friction vs support vertical).

4. Résultats Expérimentaux

Les auteurs ont évalué PRISM-WM sur des benchmarks de contrôle continu, incluant des humanoïdes de haute dimension, des configurations multi-tâches et un transfert sim-to-sim zero-shot.

• Locomotion de Haute Dimension (DiffRL & Humanoid-Bench) : PRISM-WM a atteint une efficacité d'échantillonnage et une performance asymptotique supérieures par rapport aux bases monolithiques (TD-MPC2, DreamerV3, PWM, SAC). Notamment, alors que les modèles monolithiques ont subi un effondrement de la politique dans les tâches de locomotion bipède (Humanoid, SNU-Humanoid) en raison du lissage d'événements de contact non lisses, PRISM-WM a maintenu une optimisation stable.
• Généralisation Multi-Tâches : Sur un benchmark de 30 tâches du DeepMind Control Suite, PRISM-WM a obtenu un score normalisé moyen de 0,531, soit une amélioration de 23,5 % par rapport à la base monolithique TD-MPC2 (0,430). Il a surpassé la base sur 26 tâches sur 30, atténuant efficacement le transfert négatif entre des dynamiques hétérogènes.
• Transfert Sim-to-Sim Zero-Shot : Lors du transfert d'IsaacGym vers MuJoCo, la base MLP monolithique a subi une baisse de performance de 47,8 % avec un taux de réussite de 12 %. PRISM-WM a réduit cette baisse à 11,8 % avec un taux de réussite de 88 %. L'orthogonalisation a permis au modèle d'ajuster les poids pour des experts spécifiques (par exemple, friction horizontale) sans compromettre la variété des dynamiques globales.
• Stabilité et Efficacité : Les études d'ablation ont confirmé que l'orthogonalisation prévient la dégradation de l'estimation de récompense observée dans les modèles MoE non contraints. Les tests d'efficacité d'inférence sur une NVIDIA RTX 4090 ont montré que PRISM-WM (K=4) a atteint 7 712 FPS, surpassant légèrement la base monolithique (7 499 FPS) avec une surcharge mémoire minimale.

5. Importance et Revendications

L'article revendique que PRISM-WM fournit un substrat de haute fidélité pour les algorithmes d'optimisation de trajectoire, agissant comme un modèle fondamental pour les agents basés sur un modèle. En décomposant la physique complexe en primitives orthogonales, le modèle :

• Atténue la dérive à long terme dans les horizons de planification, permettant un raisonnement fiable à long horizon dans des environnements riches en contacts.
• Stabilise la propagation des gradients pour l'apprentissage de la politique, transformant des paysages de perte fracturés en mélanges convexes fonctionnellement lisses.
• Permet une prise de décision robuste dans les systèmes hybrides où les modèles de monde standards échouent en raison d'un lissage excessif.

Les auteurs concluent que la décomposition de la physique complexe en primitives orthogonales est « fondamentalement essentielle pour maîtriser le contrôle riche en contacts ». Bien que le cadre actuel nécessite un réglage manuel du nombre d'experts $K$, ce travail établit une fondation structurelle pour construire des modèles de monde à usage général et ancrés physiquement pour les futurs agents incarnés.