One Model for All Tasks: Leveraging Efficient World Models in Multi-Task Planning

Het artikel introduceert ScaleZero, een efficiënt wereldmodel dat door middel van een Mixture-of-Experts-architectuur en een dynamische parameter-schaalstrategie (DPS) gradientconflicten oplost en met slechts 71,5% van de interacties prestaties bereikt die vergelijkbaar zijn met gespecialiseerde single-task agents voor heterogene multi-task planning.

De kern

Stel je voor dat je een superintelligente robot wilt bouwen die niet alleen één spelletje kan spelen, maar alles kan doen: van het spelen van Pong (een heel simpel spelletje met een balletje en een batje) tot het oplossen van complexe avonturen in een tekstverhaal, en zelfs het besturen van een robotarm die een kopje koffie moet pakken.

Het probleem is dat als je deze robot één brein geeft voor al deze taken, het brein vaak in de war raakt. Het is alsof je iemand probeert te leren zwemmen, piano spelen en wiskunde doen, allemaal tegelijk, door ze in één kamer te zetten. De robot wordt dan goed in het zwemmen, maar vergeet hoe hij piano moet spelen, of hij raakt vast in de wiskunde. In de wetenschap noemen we dit "plasticiteit-collaps": het brein stopt met leren omdat het te veel conflicten heeft.

De auteurs van dit paper hebben een oplossing bedacht genaamd ScaleZero. Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het Brein met "Specialisten" (De MoE Architectuur)

Stel je een groot kantoor voor. In een oud kantoor (de oude modellen) werkt iedereen in één grote open ruimte. Als de manager (de robot) een taak geeft, proberen iedereen mee te werken. Als de ene taak "rekenen" is en de andere "tekenen", schreeuwt de rekenaar tegen de tekenaar: "Nee, ik moet dit doen!" en vice versa. Het resultaat is chaos en niemand wordt goed.

ScaleZero verandert dit kantoor in een gebouw met specialisten.

• Er is een router (een slimme secretaresse) die kijkt naar de taak.
• Als het een rekenopdracht is, stuurt ze het naar het kantoor met de wiskundigen.
• Als het een tekenopdracht is, stuurt ze het naar de kunstenaars.
• Ze werken niet allemaal tegelijk aan alles. Ze werken alleen aan wat ze goed kunnen.

Dit heet in de tech-taal een Mixture-of-Experts (MoE) architectuur. Het zorgt ervoor dat de robot niet in de war raakt en blijft leren, zelfs als de taken heel verschillend zijn.

2. De Slimme Budgetplanner (Dynamic Parameter Scaling)

Nu hebben we een kantoor met specialisten, maar we hebben ook een probleem: geld en tijd.
Stel je voor dat je een robot traint om 18 verschillende taken te leren. De robot is al heel goed in "taak 1" (bijvoorbeeld een bal in een kuip vangen), maar worstelt nog met "taak 18" (een cheeta laten rennen).

In een normaal systeem zou de robot blijven oefenen aan taak 1, ook al is hij al perfect. Dat is zonde van de tijd.
ScaleZero gebruikt een slimme strategie genaamd Dynamic Parameter Scaling (DPS):

• De "Gereed" Check: Zodra de robot taak 1 onder de knie heeft, stopt hij daar direct mee. Hij "bevriest" die kennis (alsof je een boek dichtdoet en op de plank legt).
• Nieuwe Hulp: Voor de moeilijke taken die hij nog niet kan, haalt hij extra hulp bij. Hij voegt kleine, flexibele modules toe (zoals LoRA-adapters, denk aan extra gereedschap of een nieuwe set gereedschappen) die specifiek zijn voor die moeilijke taak.
• Dynamisch Groeien: Het systeem groeit mee met de moeilijkheid. Als er een nieuwe, heel moeilijke taak komt, krijgt de robot extra "hersencellen" toegewezen, maar alleen voor die taak.

Dit is alsof je een student bent die eerst zijn basiswiskunde haalt. Zodra hij dat kan, stopt hij met oefenen op basiswiskunde en begint hij met de moeilijke natuurkunde. Hij gebruikt zijn tijd veel efficiënter.

Wat is het resultaat?

De robot ScaleZero is getest op heel veel verschillende dingen:

• Atari-spelletjes (oude videospelletjes).
• DMC (robotarmen en -beesten laten bewegen).
• Jericho (tekstavonturen waar je moet lezen en beslissingen moet nemen).

De resultaten:

1. Één model voor alles: De robot doet het net zo goed als 26 verschillende robots die elk alleen één spelletje hebben geoefend.
2. Efficiëntie: Door de slimme "budgetplanner" (DPS) te gebruiken, heeft de robot 28,5% minder oefentijd nodig om even goed te worden. Hij verspilde geen tijd aan taken die hij al kon.

Samenvattend

Stel je voor dat je een meesterkok bent.

• Oude methode: Je probeert elke dag 100 verschillende gerechten te koken tegelijk in één pan. Het wordt een rommeltje en niets smaakt goed.
• ScaleZero methode: Je hebt een keuken met speciale kookstations (MoE). Als je een vis moet bakken, ga je naar het visstation. Als je een taart moet bakken, ga je naar het bakstation.
• De slimme toevoeging (DPS): Zodra je weet hoe je een taart moet bakken, stop je met oefenen op taarten en haal je een nieuw, speciaal station voor het bakken van brood. Je verspil je tijd niet aan dingen die je al kunt.

Dit paper laat zien dat we met deze slimme architectuur echt algemene robots kunnen bouwen die alles kunnen leren, zonder dat ze gek worden van de hoeveelheid informatie.

Technische samenvatting

Probleemstelling

In heterogene multi-task beslissingsproblemen (MTRL) variëren taken niet alleen in hun observatie- en actieruimtes, maar ook aanzienlijk in hun onderliggende complexiteit. Bestaande geünificeerde wereldmodellen, zoals UniZero, presteren goed in single-task settings, maar ondervinden ernstige beperkingen bij het hanteren van een brede, diverse reeks taken. De auteurs identificeren twee fundamentele obstakels:

1. Plasticiteitsverval (Plasticity Collapse): In een gedeeld model leiden conflicterende gradienten van eenvoudigere, sneller convergerende taken tot een onderdrukking van leersignalen voor complexere taken. Dit resulteert in representatieve interferentie en een progressief verlies van de plasticiteit van het netwerk (het vermogen om zich aan te passen aan nieuwe data). Dit wordt empirisch waargenomen als een toename van het percentage "dormant neurons" (inactieve neuronen) en een ongecontroleerde inflatie van de norm van de latente staat, wat leidt tot catastrofale prestatiedalingen op complexe taken.
2. Statische Resource Allocatie: Traditionele architecturen gebruiken een "one-size-fits-all" aanpak waarbij dezelfde leermiddelen (datacollectie en modelupdates) worden toegepast op alle taken, ongeacht hun intrinsieke moeilijkheidsgraad. Dit leidt tot computerefficiëntieproblemen, omdat resources worden verspild aan taken die al zijn opgelost of weinig meer opleveren, in plaats van te worden gericht op taken die nog leren vereisen.

Methodologie

De auteurs lossen deze uitdagingen op vanuit twee complementaire perspectieven: een interne architecturale oplossing en een externe procedurele strategie.

1. Architecturale Oplossing: ScaleZero

Om gradientconflicten en plasticiteitsverval te mitigeren, onderzochten de auteurs systematisch de ontwerpruimte van UniZero over vijf assen (task conditioning, encoder-architectuur, latent normalisatie, backbone-design en optimalisatiestrategie). De kern van hun oplossing is ScaleZero, een nieuw wereldmodel met de volgende kenmerken:

• Mixture-of-Experts (MoE) Backbone: In plaats van een dichte Transformer-backbone, gebruikt ScaleZero een Sparse MoE-architectuur. Hierbij worden de feed-forward netwerken vervangen door een set van gespecialiseerde "expert"-netwerken. Een router (gating mechanisme) selecteert voor elke input een subset van experts (bijv. top-1).
  • Werkingsprincipe: Taken worden via specifieke paden verwerkt door gespecialiseerde sub-netwerken. Dit reduceert gradientconflicten aanzienlijk omdat updates voor taak A niet noodzakelijk de parameters van taak B beïnvloeden.
  • Theoretische onderbouwing: De auteurs tonen aan dat de bovengrens voor gradientconflict in een MoE-laag strikt lager is dan die van een dichte laag, mits de router effectief specialisatie bevordert.
• ViT Encoder: Voor visuele taken (Atari) wordt een Vision Transformer (ViT) gebruikt, getraind vanaf nul, in plaats van een ResNet-achtige encoder.
• Layer Normalization: In plaats van SimNorm wordt standaard LayerNorm gebruikt om een betere balans te vinden tussen stabiliteit en representatieve expressiviteit.

2. Procedurele Strategie: Dynamic Parameter Scaling (DPS)

Om het probleem van statische resource allocatie op te lossen, introduceren de auteurs Dynamic Parameter Scaling (DPS). Dit is een online mechanisme dat modelcapaciteit dynamisch koppelt aan leervordering.

• Adaptieve Taak Curatie: Taken worden continu gemonitord. Zodra een taak een vooraf gedefinieerde prestatiedrempel bereikt ("solved"), wordt deze uit de actieve trainingsset verwijderd. Resources worden dan uitsluitend ingezet voor nog niet-opgeloste taken.
• Gestage Capaciteitsuitbreiding: Het trainingproces verloopt in fasen.
  • Fase 0 (Warm-up): Het basismodel wordt getraind op alle taken om algemene representaties te leren.
  • Fase s ≥ 1 (Expansion): Wanneer een nieuwe fase start (geïnitieerd door voortgang of een tijdslimiet), worden nieuwe, onafhankelijke LoRA-adapters (Low-Rank Adaptation) toegevoegd.
• Parameter Isolatie: Bij het betreden van een nieuwe fase worden de parameters van het basismodel en eerdere adapters bevroren. Alleen de nieuwe LoRA-module en schalingsfactoren worden geoptimaliseerd. Dit voorkomt catastrofale vergetelheid en zorgt voor gerichte plasticiteit voor moeilijke taken, terwijl eerder geleerde kennis behouden blijft.

Belangrijkste Resultaten

ScaleZero werd geëvalueerd op drie diverse benchmarks: Atari (26 visuele games), DeepMind Control Suite (DMC) (18 continue control-taken) en Jericho (4 tekstgebaseerde avonturenspellen).

1. Prestaties:

   • ScaleZero (als enkelvoudig multi-task agent) presteert op gelijke hoogte met of beter dan een collectie van gespecialiseerde single-task agents (UniZero per taak).
   • Op de Atari-benchmark bereikt ScaleZero een hogere gemiddelde genormaliseerde score dan 26 individueel getrainde modellen.
   • Op DMC behaalt het een superieure mediaan score, wat wijst op robuuste generalisatie.
   • Op Jericho is het prestatieniveau vergelijkbaar met state-of-the-art single-task agents en concurrerend met taak-specifieke taalmodellen.

2. Efficiëntie (DPS):

   • Wanneer ScaleZero wordt aangevuld met de DPS-strategie, bereikt het vergelijkbare prestaties als de single-task baselines maar met 28,5% minder interacties met de omgeving (op de DMC-benchmark).
   • Dit betekent een aanzienlijke reductie in data-vereisten en trainingskosten.

3. Diagnose van Plasticiteitsverval:

   • Experimenten bevestigen dat de baseline (UniZero) faalt door een toename van dormant neurons en een inflatie van de latent state norm. ScaleZero met MoE houdt deze metrics stabiel, wat de onderliggende oorzaak van het falen van eerdere modellen bevestigt.

Bijdragen

De belangrijkste bijdragen van dit werk zijn:

1. Kwantitatieve Diagnose: Een gedetailleerde analyse van plasticiteitsverval in geünificeerde wereldmodellen, waarbij een direct verband wordt gelegd tussen prestatiedaling en interne leerdynamiek (dormant neurons, latent norm).
2. ScaleZero Architectuur: De ontwikkeling van een unificerend wereldmodel dat gebruikmaakt van een MoE-backbone om gradientconflicten en representatieve interferentie effectief te mitigeren, wat leidt tot superieure prestaties op 48 diverse taken.
3. Dynamic Parameter Scaling (DPS): Een innovatieve trainingsstrategie die modelcapaciteit en resources dynamisch toewijst via LoRA-adapters, waardoor de sample-efficiency aanzienlijk wordt verbeterd zonder in te leveren op de uiteindelijke prestaties.

Significantie

Dit paper markeert een belangrijke stap in de ontwikkeling van generalist agents voor reinforcement learning. Het toont aan dat het mogelijk is om één enkel model te trainen dat complexe, heterogene taken kan beheersen zonder in te leveren op prestaties, mits de juiste architecturale (MoE) en procedurele (DPS) maatregelen worden genomen. De combinatie van conditional computation en adaptieve resource allocatie biedt een robuust fundament voor toekomstige systemen die moeten opereren in dynamische, multi-modal omgevingen. De code is beschikbaar gesteld via GitHub, wat de reproduceerbaarheid en verdere onderzoek in dit domein faciliteert.