Technische Samenvatting: Prisma Wereldmodel (PRISM-WM)

1. Probleemstelling

Modelgebaseerde Versterkende Leer (MBRL) en planning in robotische domeinen worden fundamenteel beperkt door het hybride karakter van fysieke dynamica. Wereldse systemen, met name bij locomotie op poten en manipulatie, omvatten continue beweging die wordt onderbroken door discrete gebeurtenissen (bijv. contacten, impacten, vastzitten versus glijden). Deze interacties creëren niet-gladde, sterk niet-lineaire dynamica met scherpe discontinuïteiten in het onderliggende vectorveld.

Conventionele latente wereldmodellen maken doorgaans gebruik van monolithische neurale netwerken die globale continuïteit afdwingen. Hoewel deze modellen theoretisch in staat zijn tot universele benadering, leren ze in de praktijk een gemiddelde representatie die oververmijdt onderscheidende dynamische modi. Voor planners voor trajectoptimalisatie veroorzaakt deze oververmijding cumulatieve fouten tijdens voorspellingen met een lange horizon, met name bij fysieke grenzen. Dit resulteert in:

• Spectrale vervaging over contactgrenzen heen, wat leidt tot aanzienlijke gradiëntmismatch.
• Cumulatieve drift in hoogdimensionale configuratieruimten, waardoor langetermijnplanning onbetrouwbaar wordt.
• Beleidinstorting in contactrijke omgevingen waar standaardmodellen niet kunnen onderscheiden tussen stabiele en instabiele regimes.

2. Methodologie: Het Prisma Wereldmodel (PRISM-WM)

Om deze beperkingen aan te pakken, introduceren de auteurs PRISM-WM, een gestructureerde architectuur die is ontworpen om complexe hybride dynamica te ontleden in samenstelbare primitieven. Het model is gebaseerd op een contextbewust Mixture-of-Experts (MoE)-kader met een nieuw latent orthogonaal doelwit.

Kernarchitectuur

• Contextbewuste Gating: Een gating-netwerk ($G$) ontvangt een contextvector $c_t$ (bestaande uit toestand $z_t$, actie $a_t$, en optioneel taakembedding $e_{task}$) en produceert softmax-genormaliseerde gewichten $w_t$. Dit mechanisme identificeert impliciet het actieve fysieke regime (bijv. vlucht versus stand) zonder expliciete supervisie.
• Gespecialiseerde Experts: De toestand wordt doorgestuurd naar een set van $K$ expertnetwerken $\{E_k\}$. Elke expert is gespecialiseerd in het leren van de continue stroming binnen een specifiek lokaal dynamisch regime, waardoor interferentie van conflicterende modi wordt vermeden.
• Residupredictie: De volgende latente toestand wordt voorspeld als een gewogen combinatie van expertuitkomsten die worden opgeteld bij de huidige toestand:
  $$\hat{z}_{t+1} = z_t + \sum_{k=1}^{K} w_{t,k} \cdot v_k(x_t)$$
  waarbij $v_k$ de residudynamica voorstelt die wordt voorspeld door de $k$-de expert.

Belangrijkste Innovatie: Latente Orthogonalisatie

Om modinstorting te voorkomen (waarbij experts redundante kenmerken leren) en diversiteit te waarborgen, introduceert PRISM-WM een differentieerbare Gram-Schmidt-orthogonalisatielaag.

• Nadat elke expert een vector $u_k$ van de voorlaatste laag heeft gegenereerd, past het model een gestabiliseerd Gram-Schmidt-proces toe om orthogonale basisvectoren $v_k$ te genereren.
• Dit dwingt de expertrepresentaties om op een Stiefel-mannigvuldigheid te liggen, waardoor wordt gegarandeerd dat ze een subruimte met maximale rang en zonder redundantie beslaan voor de residudynamica.
• Hoewel fysieke contactmodi ongeordend zijn, zorgt deze orthogonalisatie ervoor dat de experts verschillende componenten van de residumannigvuldigheid vastleggen, waardoor het gating-netwerk ongeordende fysieke modi kan reconstrueren via een gewogen som.

Integratie met Planning

PRISM-WM is geïntegreerd in het TD-MPC-planningskader. Door complexe dynamica te factoriseren in niet-redundante experts, reduceert het model het residu van het model in één stap van $\epsilon$ tot ongeveer $\epsilon/K\alpha$. Deze structurele ontleding mitigeert cumulatieve drift, waardoor hoogtrouwrollouts met een lange horizon en stabiele gradiëntpropagatie voor beleidsoptimalisatie mogelijk worden.

3. Belangrijkste Bijdragen

1. Structurele Ontleding: Introductie van PRISM-WM, dat hybride dynamica ontleden met een contextbewuste MoE-architectuur gekoppeld aan latente orthogonalisatie, en zich afwendt van monolithische globale benaderingen.
2. Orthogonalisatie-doelwit: Een nieuw trainingsdoelwit met behulp van differentieerbare Gram-Schmidt-orthogonalisatie om expertdiversiteit af te dwingen en modinstorting te voorkomen, waardoor de latente ruimte structureel wordt afgestemd op discrete fysieke overgangen.
3. Theoretische en Empirische Validatie: Aantonen dat deze ontleding cumulatieve fouten in hoogdimensionale ruimten reduceert en de geometrie van het verlieslandschap verbetert voor zowel planning (TD-MPC) als direct beleidslernen (PWM).
4. Zero-Shot Transfermogelijkheid: Bewijs dat de orthogonale basis het model in staat stelt om te generaliseren over simulatoren met verschillende contactfysica (bijv. zacht-contact IsaacGym naar stijf-contact MuJoCo) door specifieke dynamische modi te isoleren (bijv. wrijving versus verticale ondersteuning).

4. Experimentele Resultaten

De auteurs hebben PRISM-WM geëvalueerd op benchmarks voor continue besturing, waaronder hoogdimensionale humanoïden, multi-taakinstellingen en zero-shot sim-naar-sim transfer.

• Hoogdimensionale Locomotie (DiffRL & Humanoid-Bench): PRISM-WM behaalde een hogere sample-efficiëntie en asymptotische prestatie in vergelijking met monolithische baselines (TD-MPC2, DreamerV3, PWM, SAC). Opvallend is dat, terwijl monolithische modellen beleidsinstorting ondervonden in bipedale locomotietaken (Humanoid, SNU-Humanoid) door het middelen van niet-gladde contactgebeurtenissen, PRISM-WM stabiele optimalisatie handhaafde.
• Multi-taak Generalisatie: Op een DeepMind Control Suite-benchmark met 30 taken behaalde PRISM-WM een gemiddelde genormaliseerde score van 0,531, een 23,5% verbetering ten opzichte van de monolithische TD-MPC2-baseline (0,430). Het presteerde beter dan de baseline op 26 van de 30 taken, waardoor negatieve transfer tussen heterogene dynamica effectief werd gemitigeerd.
• Zero-Shot Sim-naar-Sim Transfer: Bij overdracht van IsaacGym naar MuJoCo leed de monolithische MLP-baseline aan een 47,8% daling in prestaties met een slagingspercentage van 12%. PRISM-WM verlaagde deze daling tot 11,8% met een slagingspercentage van 88%. De orthogonalisatie stelde het model in staat om gewichten voor specifieke experts (bijv. horizontale wrijving) aan te passen zonder de globale dynamica-mannigvuldigheid te compromitteren.
• Stabiliteit en Efficiëntie: Ablatiestudies bevestigden dat orthogonalisatie degradatie van de beloningsschatting voorkomt zoals gezien in onbeperkte MoE-modellen. Tests voor inferentie-efficiëntie op een NVIDIA RTX 4090 toonden aan dat PRISM-WM (K=4) 7.712 FPS behaalde, iets beter dan de monolithische baseline (7.499 FPS) met een minimaal geheugenoverhead.

5. Betekenis en Claims

Het artikel beweert dat PRISM-WM een hoogtrouw substraat biedt voor trajectoptimalisatie-algoritmen, dienend als een fundamenteel model voor modelgebaseerde agenten. Door complexe fysica te ontleden in orthogonale primitieven, doet het model het volgende:

• Het mitigeert langetermijndrift in planningshorizons, waardoor betrouwbare redenering met een lange horizon mogelijk wordt in contactrijke omgevingen.
• Het stabiliseert gradiëntpropagatie voor beleidsoptimalisatie, waardoor gebroken verlieslandschappen worden omgezet in functioneel gladde convexe mengsels.
• Het maakt robuuste besluitvorming mogelijk in hybride systemen waar standaard wereldmodellen falen door oververmijding.

De auteurs concluderen dat het ontleden van complexe fysica in orthogonale primitieven "fundamenteel essentieel is voor het beheersen van contactrijke besturing". Hoewel het huidige kader handmatige afstemming van het aantal experts $K$ vereist, legt het werk een structurele basis voor het bouwen van algemeen toepasbare, fysisch onderbouwde wereldmodellen voor toekomstige embodied agenten.