Een wereldmodel leren om Pokémon te spelen

San Francisco was onlangs overspoeld met Pokémon-memorabilia: grote advertenties met vele van de 1026 pocketmonsters waren overal op de metrolijnen en bushaltes in Market Street geplakt. SF was dit jaar de locatie van de Pokémon World Championships, en veel eclectische en vreugdevolle Pokémon-trainers waren te vinden in de heuvels en het financiële district van downtown SF. Misschien hadden alle advertenties me onbewust beïnvloed, maar ik besloot dat een goede "wereld" voor ons wereldmodel het spel uit 1996 zou zijn waar het allemaal mee begon: Pokémon Red.

De mechanieken van Pokémon

Pokémon Red is een spel waarin je een wereld verkent, wezens genaamd Pokémon verzamelt en deze in gevechten gebruikt. Je gebruikt de richtingstoetsen om rond te lopen of door menu's te navigeren. De A-knop wordt gebruikt om met objecten te communiceren, dialogen verder te helpen en keuzes te bevestigen; de B-knop annuleert meestal of gaat een stap terug. Het kiezen van een starter betekent dat je naar een Poké Ball (een container voor Pokémon) toe gaat en de dialoog doorloopt die je keuze bevestigt.

Aan het begin van het spel bevindt de speler zich in het laboratorium van Professor Oak, waar hij je je eerste Pokémon aanbiedt: Bulbasaur, Charmander of Squirtle. Dit zijn de drie "starters", die elk in een Poké Ball op een tafel in zijn lab wachten.

Het Plan

Het doel was relatief eenvoudig: versla de kleinzoon van Professor Oak. Om dit in stappen te verdelen, bedacht ik het volgende:

  1. Ga naar het laboratorium van Professor Oak.
  2. Voltooi de dialoog met Oak.
  3. Selecteer een starter-Pokémon.
  4. Probeer het laboratorium te verlaten.
  5. Confronteer en versla de kleinzoon van Oak.

Al snel besefte ik dat dit wellicht te ambitieus was voor een eerste experiment. Daarom versmalde ik het doel tot het selecteren van een starter vanuit een opgeslagen staat in het laboratorium van Oak, waarbij het verkrijgen van een van de drie starters als een succesvolle poging zou tellen.

Vanaf de opgeslagene positie in het lab is het indrukken van de A-knop twaalf keer voldoende. Maar kan het model dat leren? Het model zou ook doelloos rond kunnen dwalen, of de B-knop kunnen indrukken na elke paar sequenties van A, waardoor het zijn inspanningen annuleert net voordat het doel bereikt is.

Wat is een wereldmodel precies?

Een wereldmodel is een model waarin we beginnen met een huidige staat of observatie, waarna er een actie plaatsvindt die deze staat wijzigt en een nieuwe observatie produceert. Realistisch gezien moet er een correlatie (of idealiter causaliteit) zijn tussen de actie in een bepaalde staat en de nieuwe staat die daaruit voortvloeit. Het doel van het wereldmodel is om deze correlatie te leren.

Stel dat de huidige observatie een screenshot is van het laboratorium van Professor Oak, met de speler voor de Poké Balls. De actie is het indrukken van "links", waardoor de speler één tegel naar links beweegt, wat resulteert in een nieuw screenshot. Het doel is dat het wereldmodel een vorm van intuïtie ontwikkelt dat "links" indrukken de speler naar links verplaatst.

$$o{t+1}\approx F(ot,a_t)$$

Hier is $ot$ het huidige screenshot, $at$ is de ingedrukte knop, en $F$ is de functie die het wereldmodel leert om het volgende screenshot $o_{t+1}$ te voorspellen. Het doel is om te generaliseren: het model moet niet alleen leren dat links indrukken in het lab van Professor Oak de speler naar links verplaatst, maar dat links indrukken overal de speler naar links zou moeten verplaatsen.

Terzijde: deze opzet lijkt misschien op reinforcement learning. Maar cruciaal is dat het wereldmodel de dynamiek van de staat leert begrijpen zonder een beloning, oftewel: het is reward-free.

Hoe het wereldmodel leert voorspellen

Van screenshots naar embeddings

Als we echt geïnteresseerd zijn in de staatsovergangen en de dynamiek van de omgeving, en we hebben alleen observaties (in dit geval screenshots), leert het model dan screenshots te voorspellen?

Nee, wat het model feitelijk leert, is voorspellen binnen zijn latente ruimte, ook wel de embedding space genoemd. We gebruiken eerst een encoder die, wanneer hij een screenshot krijgt, een embedding produceert.

Wat is een embedding? Een embedding is een geleerde representatie van een input als een vector van getallen. Een encoder kan een afbeelding, een zin of een geluid omzetten in zo'n vector, waardoor een ander model iets heeft dat het kan vergelijken, voorspellen of als input kan gebruiken. De waarden zijn geen handmatig gelabelde kenmerken; één coördinaat hoeft niet direct "kleur" of "positie" te betekenen. Informatie kan over vele vermeldingen verspreid zijn, en sommige details van de oorspronkelijke input kunnen volledig verdwijnen.

Met de embedding die de huidige observatie $ot$ representeert, kunnen we een predictor-model gebruiken dat, gegeven een actie $at$, probeert de toekomstige embedding te voorspellen. Deze voorspelde embedding kunnen we vervolgens vergelijken met de werkelijke embedding die wordt geproduceerd door de toekomstige observatie $o_{t+1}$. Dit is relatief simpel: we sturen het nieuwe screenshot door dezelfde encoder en krijgen de werkelijke toekomstige embedding. Het doel is dan:

$$\hat z{t+1}\approx z{t+1}$$

Hier is $\hat z{t+1}$ de voorspelde volgende embedding en $z{t+1}$ de embedding van het screenshot dat daadwerkelijk volgde. Uitgeschreven voor één stap:

$$zt = E(xt), \qquad \hat z{t+1} = P(zt,at), \qquad z{t+1} = E(x_{t+1})$$

Hierbij is $xt$ het huidige screenshot, $at$ de actie, en $E$ en $P$ de encoder en predictor. Belangrijk is dat de voorspelde embedding $\hat z{t+1}$ nooit wordt teruggezet naar een screenshot; we vergelijken het direct met $z{t+1}$.

Voorspellingsverlies en collaps

Om het verschil tussen onze voorspelling en de werkelijke embedding te leren, gebruiken we de Mean-Squared Error (MSE) verliesfunctie:

$$\mathcal L{\mathrm{pred}}=\frac{1}{D}\sum{d=1}^{D}(\hat z{t+1,d}-z{t+1,d})^2,\qquad D=192$$

Dit is in essentie een afstandsfunctie; we meten de afstand tussen onze twee embedding-vectoren.

Er is echter een gevaarlijk degeneratief scenario mogelijk. Omdat we zowel de encoder als de predictor trainen en de encoder twee keer voorkomt in elk trainingsvoorbeeld, kan het volgende gebeuren: de encoder leert elk screenshot als dezelfde embedding $c$ te representeren. De predictor leert vervolgens simpelweg dat de volgende embedding ook $c$ zal zijn. Dit wordt latent collapse genoemd. Het model heeft een triviale manier gevonden om de verliesfunctie te minimaliseren door alle embeddings tot één enkel punt te reduceren.

SIGReg: de embedding-ruimte bruikbaar houden

Om te voorkomen dat embeddings inklappen tot één punt, is een vorm van regularisatie nodig. De LeWorldModel-paper kiest voor een relatief eenvoudige oplossing genaamd SIGReg (Sketched Isotropic Gaussian Regularization).

SIGReg moedigt de embeddings aan om over een batch screenshots te lijken op een standaard isotrope Gaussiaan:

$$z\sim\mathcal N(\mu,\Sigma),\qquad \mu=\mathbf 0,\quad \Sigma=I_D$$

Een isotrope Gaussiaan heeft een gemiddelde van 0 en de identiteitsmatrix $I_D$ als covariantiematrix. In twee dimensies is dit een ronde wolk in plaats van een punt of een lijn.

SIGReg kiest een willekeurige richting $u$ met lengte één en meet waar elke embedding langs deze richting valt:

$$hi=u^\top zi,\qquad i=1,\ldots,N$$

Als de oorspronkelijke wolk echt $\mathcal N(0,I_D)$ volgde, zouden de getallen langs elke eenheidsrichting $\mathcal N(0,1)$ volgen. SIGReg beantwoordt de vraag of deze geprojecteerde getallen lijken op monsters van een standaard Gaussiaan met een karakteristieke-functietest.

De Karakteristieke Functie Een karakteristieke functie is een andere manier om de vorm van een kansverdeling te beschrijven. Voor een willekeurige vector $X$ is deze gedefinieerd als:

$$\phi_X(t)=\mathbb{E}[e^{i t^\top X}]$$

Voor een isotrope Gaussiaan $X\sim\mathcal N(0,\sigma^2 I)$ is de karakteristieke functie:

$$\phi_X(t)=e^{-\frac12\sigma^2\|t\|^2}$$

Voor elke richting vergelijkt SIGReg de geschatte karakteristieke functie uit de geprojecteerde monsters met de bekende functie voor $\mathcal N(0,1)$. De Epps-Pulley-test wordt vervolgens gebruikt om dit verschil om te zetten in een straf (penalty). Dit gemiddelde over 1.024 willekeurige richtingen wordt gebruikt om de encoder bij te werken.

De uiteindelijke objectieffunctie voor de training ziet er als volgt uit:

$$\mathcal L=\mathcal L_{\mathrm{pred}}+0.1\,\operatorname{SIGReg}(Z)$$

De factor 0.1 bepaalt hoeveel de regularisator bijdraagt om het wissen van variaties tussen screenshots tegen te gaan.

De Pokémon Trainingsdata

Ik heb 42.382 grijswaardenframes opgenomen vanuit de Pokémon Red-emulator, verdeeld over 1.009 korte trajecten. Sommige trajecten volgen gescripte routes naar een starter; andere voegen ruis toe aan die routes of bewegen meer willekeurig. Voor elke stap bevat de opname het huidige frame, de ingedrukte knop en het daaropvolgende frame.

Het is essentieel om "rommelige" trajecten te hebben. Een model dat alleen op een schone route is getraind, zou bijvoorbeeld nooit leren wat de B-knop doet tijdens een dialoog. De planner zal echter allerlei knopsequenties voorstellen, inclusief foutieve, en heeft voorspellingen nodig voor die sequenties.

Voordat ik begon met plannen, wilde ik weten of de embeddings bijhielden of de speler een Pokémon had gekozen. Ik bevroor de encoder en trainde een kleine classifier (een linear probe) om dit te bepalen op basis van de embeddings. De classifier kon inderdaad bepalen of de speler een Pokémon had, zelfs op beelden die niet in de training zaten. Ook de predictor presteerde beter dan een baseline die simpelweg de huidige embedding kopieerde.

Plannen in de geleerde wereld

Om een starter te vragen, nam ik screenshots van succesvolle selecties van Bulbasaur, Charmander en Squirtle. Hun embeddings vormen een set van drie mogelijke doelen, $\mathcal G$. Dit geeft de planner een beeld van het gewenste resultaat zonder dat er een route of beloningen per knopdruk worden meegegeven.

Het proces werkt als volgt: we beginnen met de werkelijke screenshot van de opgeslagen staat, encoderen deze, en geven die embedding samen met de eerste voorgestelde knop aan de predictor. Voor de tweede knop is er geen nieuw screenshot; we moeten de eerste voorspelling terugvoeren. Een kandidaat-plan ontvouwt zich zo:

$$\hat z0=E(x0),\qquad \hat z{t+1}=P(\hat zt,a_t)$$

We berekenen een score $J$ voor elke voorspelde sequentie van veertien knoppen:

$$J(a{0:H-1})=\min{1\leq k\leq H}\;\min{zg\in\mathcal G}\frac{1}{D}\|\hat zk-zg\|_2^2,\qquad H=14$$

In simpele woorden: na elke denkbeeldige knopdruk vergelijken we de voorspelde embedding met de drie starter-doelen en houden we de kleinste afstand bij. De uiteindelijke score $J$ is de kleinste afstand die we over de hele reeks van veertien stappen hebben gezien.

Omdat er te veel combinaties zijn om één voor één te proberen, gebruikte ik de cross-entropy method (CEM) om het zoekgebied te verkleinen:

  1. Sample 512 complete plannen en gebruik het wereldmodel om te voorspellen wat ze doen.
  2. Bereken de $J$-score voor elk plan en behoud de 64 plannen met de laagste scores.
  3. Pas de waarschijnlijkheid van de knoppen in elk slot aan, zodat de knoppen uit de beste plannen vaker worden gesampled in de volgende ronde.

Dit proces wordt herhaald over meerdere rondes, waarna het beste plan in de emulator wordt uitgevoerd.

Waarom het eerste plan faalde

De eerste zoekopdracht vond een plan waarvan het model dacht dat het dicht bij een starter zou komen, maar in de emulator leidde dit niet tot het verkrijgen van een Pokémon.

Een waarschijnlijk probleem was de training. Tijdens de oorspronkelijke training begon de predictor elke stap met een embedding van een echt screenshot. Bij het plannen is er echter geen reset: de tweede voorspelling begint bij de eerste voorspelling, enzovoort. Kleine fouten kunnen zich opstapelen totdat het model werkt met embeddings die niet meer lijken op echte screenshots.

CEM kan dit verergeren. Het zoekt naar sequenties die volgens het model het beste zijn. Als het model een grote fout maakt bij een specifieke sequentie, kan die fout ervoor zorgen dat de sequentie onterecht zeer aantrekkelijk lijkt voor de zoekfunctie.

Rollout Fine-Tuning en de tweede poging

Om de predictor te laten oefenen met zijn eigen fouten, heb ik hem gefinetuned op rollouts. Een rollout begint bij een echt screenshot, maar de predictor gebruikt vervolgens zijn eigen voorspellingen als basis voor de volgende stappen, in plaats van de werkelijke volgende screenshots. We gebruiken de echte screenshots nog wel als doel om de voorspellingen te controleren.

Na deze fine-tuning werd de eerste voorspelling paradoxaal genoeg iets slechter, maar de foutmarge groeide veel langzamer over langere sequenties. Dit is de tradeoff die belangrijk is: een plan moet standhouden over meer dan één knopdruk.

Bij de tweede poging vond de planner een sequentie die Squirtle selecteerde. Hoewel dit niet de meest directe route was (zoals twaalf keer A indrukken), veranderde de party-count in de emulator van nul naar één. De voorspelde successen kwamen nu overeen met de realiteit in het spel.

Om te testen of dit toeval was, draaide ik de planner met 100 verschillende willekeurige seeds. 52 van de 100 plannen wisten een starter te bemachtigen, vergeleken met nul bij volledig willekeurige knopsequenties en slechts één bij een niet-getrainde predictor.

Het einde…?

Ik heb de kleinzoon van Oak (Blue of Gary) nooit verslagen, maar het model kan in ieder geval iets. Het uiteindelijke model, getraind vanaf nul, bevat ongeveer 12,5 miljoen parameters. Ik denk dat er interessantere versies mogelijk zijn, zoals het model laten starten buiten het lab, hem naar Oak laten lopen, de volledige dialoog laten doorlopen en dan pas een starter laten kiezen. Ik vermoed echter dat de moeilijkheidsgraad exponentieel toeneemt met de lengte van het plan.

Voor nu heb ik in ieder geval Squirtle.

***

Referenties Lucas Maes, Quentin Le Lidec, Damien Scieur, Yann LeCun, and Randall Balestriero. LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. 2026. arXiv:2603.19312.