Het project Tiny World Former demonstreert hoe twee kleine neurale netwerken kunnen worden ingezet op een Raspberry Pi Pico 2 W om games te simuleren en te spelen. Het systeem bestaat uit een wereldmodel (een looped transformer) dat het volgende frame en de beloning voorspelt, en een actor-netwerk (een MLP) dat acties kiest op basis van de gegenereerde frames.
Belangrijke technische aspecten zijn:
- Games: Er zijn vijf versies geïmplementeerd: Breakout, Doom-lite, Flappy Bird, Pong en Pac-Man lite.
- Optimalisatie: Om op de microcontroller te passen, is gebruikgemaakt van
int8 quantisatie en Quantization Aware Training (QAT). De firmware is geschreven in no_std Rust op basis van Embassy.
- Training: De agents zijn getraind met Stream AC(λ), een methode die leert van individuele transities zonder replay buffer.
- Hardware: De RP2350-chip gebruikt core 1 voor de game loop en core 0 voor USB-communicatie en displaybeheer.
Naast de fysieke hardware is er ook een WebAssembly-versie beschikbaar die de volledige loop in een browser simuleert.
Tiny World Former: Games genereren en spelen op een Raspberry Pi Pico 2 W
Alle frames worden gesampled door het wereldmodel-netwerk, en de acties worden uitgevoerd door het actor-netwerk.
Motivatie
De aanleiding voor dit project was interesse in looped transformers, waarbij één transformer-blok meerdere keren wordt toegepast met dezelfde gewichten in plaats van het stapelen van verschillende lagen. Omdat het hergebruiken van gewichten weinig geheugen kost, was de vraag of een looped world model klein genoeg zou zijn om op een microcontroller te draaien. De keuze viel op de Pico 2 W vanwege de lage prijs en de aanwezigheid van een FPU (Floating Point Unit). Er is gekozen voor eenvoudige games, omdat de echte versies direct in numpy geschreven kunnen worden om elk gegenereerd frame te controleren. Later is een streaming RL-agent toegevoegd, zodat het bord ook kan leren spelen binnen zijn eigen "droom".
Het wereldmodel en de actor
Overzicht van de netwerken
| Kenmerk | Netwerk 1: Het Wereldmodel | Netwerk 2: De Actor |
|---|
| Taak | frame + actie → volgend frame + reward | frame → actie |
| Architectuur | looped transformer, q int8 | MLP met twee verborgen lagen, float32 |
| Grootte | 26k–131k parameters | 9k–16k parameters |
| Kosten per frame | 2–8M multiply-adds | 9k–16k multiply-adds |
| Training | getraind op opgenomen echte games | met stream-x, in het echte spel of in de droom |
De grootte van het model hangt af van het spel; sommige games vereisen meer parameters om speelbaar te worden. Er is gestreefd naar een zo klein mogelijke parameterset.
Op de chip wordt het echte spel door geen van beide netwerken gebruikt: de actor ziet alleen frames van het wereldmodel, en het wereldmodel ontvangt alleen acties van de actor. Beiden draaien op core 1 van de RP2350, terwijl core 0 de USB en het display beheert.
De vooraf getrainde actor kan worden vervangen door een nieuwe stream-x agent die vanaf nul op het bord leert. Bovendien kan de gebruiker via een knopdruk de controle overnemen en zelf in de "droom" spelen.
De spellen
Elk spel bestaat uit enkele honderden regels numpy in src/tinyworldformer/games/. De rendering gebeurt naar een klein raster van token-id's. Eén frame bevat altijd de volledige spelstatus, waardoor de verborgen status in het frame moet worden getekend (bijv. de richting van de bal is onderdeel van het token, en Doom-lite heeft een minimap).
| Spel | Frame | Regels | Random event | Wereldmodel (params / MACs) | Exacte frames |
|---|
| Breakout | 20×16 | paddle, bal, 56 bricks, 5 levens | de serve na een miss | 26k / 4.3M | 99.99% |
| Doom-lite | 32×32 | raycast first-person view van 8×8 doolhof; vijand loopt kortste route; schieten | respawn van de vijand | 55k / 4.2M | 83.6% |
| Flappy Bird | 24×24 | zwaartekracht, flaps, scrollende pipes | hoogte van de volgende opening | 46k / 2.0M | 100% |
| Pong | 20×32 | speler tegen CPU (reageert traag), eerste tot 5 | de serve na een punt | 46k / 2.3M | 99.45% |
| Pac-Man lite | 48×48 | 23×21 doolhof, 249 pellets, twee ghosts op kortste route, 3 levens | respawn van de ghosts | 131k / 8.2M | 86.4% |
"Exacte frames" is het aandeel van achtergehouden overgangen waarbij elke cel van het voorspelde frame exact overeenkomt met het echte spel. Frames direct na een random event zijn uitgesloten, omdat deze alleen gesampled kunnen worden en niet voorspeld. Bij Doom-lite is het percentage lager omdat één foutieve pixel in de 3D-view het hele frame als fout markeert, terwijl de spelstatus bijna altijd correct is.
De echte games kunnen in een terminal worden gespeeld via: uv run python -m tinyworldformer.games.play --game pacman (of breakout, doomlite, flappy, pong).
Hoe het wereldmodel werkt
Frames opgedeeld in tokens
Het frame wordt verdeeld in vierkante patches: 2×2 cellen voor Breakout, 4×4 voor Doom-lite, Flappy Bird en Pong, en 6×6 voor Pac-Man. Dit resulteert in 36 tot 80 tokens per spel. De embedding van een patch is de som van één geleerde vector per cel, met een aparte tabel voor elke positie binnen de patch om de celvolgorde te behouden. Elke token krijgt daarnaast een positievector en de vector van de actie.
Looped Transformer
Het model is een enkel transformer-blok (attention plus een ReLU MLP) met één set gewichten, toegepast 1 tot 4 keer, waarbij de input voor elke pass opnieuw wordt toegevoegd. De output na elke pass is getraind tegen het volgende frame. Tijdens runtime is het aantal loops een afweging tussen snelheid en nauwkeurigheid. Op de Pico worden twee loops gebruikt om een acceptabele framerate te behouden.
Exacte achtergehouden frames per aantal loops:
| Loops | Breakout | Doom-lite | Flappy Bird | Pong | Pac-Man |
|---|
| 1 | 99.68% | 59.8% | 99.99% | 98.47% | 79.6% |
| 2 | 99.99% | 83.6% | 100% | 99.45% | 86.4% |
| 3 | 100% | 85.9% | 100% | 99.60% | 85.7% |
| 4 | 99.96% | 84.7% | 100% | 99.53% | 84.4% |
Decoding
Het model produceert een score voor elke cel en elk token. Het omzetten van deze scores naar een frame gebeurt via drie regels:
- Achtergrondcellen (muren, bricks, pellets, vloer): elke cel neemt onafhankelijk het meest waarschijnlijke token.
- Singletons (de bal, de speler, elke ghost): deze worden geplaatst via één keuze over alle (cel, token) paren, zodat er altijd precies één van elk is. Na een random event wordt de keuze gesampled tussen kandidaten binnen 2.2 logits van de beste optie.
- Sprites: Pac-Man en de ghosts zijn 2×2 sprites, maar slechts één cel draagt de status (bijv. richting). De overige drie cellen worden hiervan afgeleid in plaats van voorspeld.
Validatie van het wereldmodel
Om te controleren of het model robuust is, speelt een agent in de droom. De spelstatus wordt teruggelezen uit het gedroomde frame, de echte regels worden toegepast op die status en actie, en het resultaat wordt vergeleken met het volgende gedroomde frame.
Resultaten over 5.000 gedroomde stappen (bij 2 loops):
| Spel | Foutieve spelstatussen | Foutieve rewards | Foutieve cellen per stap |
|---|
| Breakout | niet teruggelezen; 0 frames met onmogelijke tokens | 6 | – |
| Doom-lite | 1 | 0 | 1.50 |
| Flappy Bird | 0 | 0 | 0.00 |
| Pong | 1 | 0 | 0.00 |
| Pac-Man | 742 | 38 | 0.39 |
Training
- Dataverzameling: Agents spelen het echte spel; elke stap wordt opgenomen als een voorbeeld (frame, actie, volgend frame, reward). Er worden 1 tot 3 miljoen voorbeelden per spel verzameld.
- Training Wereldmodel: Het model leert het volgende frame te voorspellen met een per-cel cross-entropy loss (30k stappen). De eerste 3–4k stappen zijn in float, daarna schakelt Quantization Aware Training (QAT) in voor int8 rounding.
- Implementatie: De int8-gewichten worden in de Rust-broncode geschreven. De firmware draait vervolgens op de Pico.
- Fine-tuning: Een agent die in het wereldmodel is getraind, kan statussen bereiken die niet in de oorspronkelijke data zaten. Er worden nieuwe echte games opgenomen met deze agent, waarna het wereldmodel daarop wordt gefinetuned.
Gevonden exploits en fixes
- Dubbele objecten: Het wegen van de loss op basis van "cel verandering" zorgde ervoor dat het model twee ballen tekende. Dit is opgelost door weging op basis van waar actie mogelijk is in het inputframe en de globale singleton-keuze.
- Logische fouten: Een Breakout-agent scoorde 154 punten terwijl er 26 bricks overbleven, omdat het model bricks in de hoeken niet verwijderde bij impact.
- Droom vs. Realiteit: Agents die alleen in de droom trainden, scoorden daar 210 (max realiteit is 56) maar slechts 28 in het echte spel.
- Sprite-verlies: In Doom-lite verdween de vijand soms uit de droom. Noise-augmentatie (random pixels toevoegen aan input views) heeft dit opgelost.
De agents: Stream-X
De agents maken gebruik van Stream AC($\lambda$), wat leert van één transitie per keer zonder replay buffer, batches of target networks. Eligibility traces sturen credit terug in de tijd, en ObGD voorkomt dat stappen het TD-target overschrijden. De policy- en value-netwerken hebben elk twee verborgen lagen van 64 units (32 op de Pico) met LayerNorm.
Input features per spel:
- Breakout: rij, kolom en richting van de bal.
- Pac-Man: voor elke richting de kortste afstand tot de dichtstbijzijnde pellet en de afstand tot elke ghost.
Greedy return in het echte spel (gemiddelde van 10 episodes):
| Spel | 0 | 10k | 50k | 100k | 250k | Max |
|---|
| Breakout | 8.2 | 29.9 | 29.9 | 31.1 | 45.6 | 56 |
| Doom-lite | −5.0 | −5.0 | 69.9 | 78.1 | 82.9 | – |
| Flappy Bird | −1.0 | 24.6 | 96.2 | 123.7 | 332 | 332 |
| Pac-Man | 7.0 | 17.7 | 72.9 | 146.6 | 217.4 | 249 |
| Pong | −4.3 | −2.0 | −0.3 | 1.1 | 1.1 | 5 |
Om te voorkomen dat Doom-lite agents alleen in de hoek blijven staan schieten, is een --idle-penalty 0.1 toegevoegd.
Quantisatie en de Pico
Elke matrix in het wereldmodel is int8 met één schaal per output-rij. Activaties zijn eveneens int8. Matrixvermenigvuldigingen sommeren int8 × int8 naar int32 (gebruikmakend van de SMLAD instructie op de RP2350). Softmax maakt gebruik van een exp-tabel van 1024 entries. Norms, residuals en embeddings blijven in float32 vanwege de FPU op de RP2350.
Geheugengebruik
| Spel | int8 gewichten | float32 (emb, norms, scales) | Agent (float32) | Firmware Flash | Firmware RAM |
|---|
| Breakout | 19.7 KB | 31.6 KB | 54.1 KB | 361 KB | 299 KB |
| Doom-lite | 33.9 KB | 91.9 KB | 61.3 KB | 471 KB | 357 KB |
| Flappy Bird | 30.1 KB | 69.5 KB | 50.8 KB | 409 KB | 227 KB |
| Pong | 30.1 KB | 70.4 KB | 52.6 KB | 414 KB | 244 KB |
| Pac-Man | 131k params | - | - | 600 KB | 460 KB |
Op de chip
De firmware is geschreven in no_std Rust op Embassy. Er vindt geen runtime-allocatie plaats; gewichten staan in flash en buffers zijn vaste statische variabelen in RAM.
- Core 1: Draait de game loop. De actor kiest een actie → het int8 transformer model voert 2 loops uit → decoding zet scores om in een frame → de reward wordt bepaald.
- Core 0: Beheert async taken. Verstuurt frames via USB-serieel als gekleurde ANSI-tekst en beheert het optionele Display Pack (ST7789 over SPI).
Geheugenverdeling:
- Flash: Wereldmodel (int8 matrices, float32 embeddings), actor (MLP gewichten) en firmware code.
- RAM: Scratch buffers voor forward passes, features en hidden layers van de actor, en weights/traces voor de on-device learner.
Flashing en Uitvoering
Het flashen gebeurt via het tinyworldformer-rs crate. Voorbeelden van commando's:
./flash.sh --env doomlite --agent doomagentroam_p0.3 (bouwen en flashen)./flash.sh --env breakout --display (met Display Pack)./flash.sh --env pong --host (spelen in de terminal)
Besturing via USB:
- Breakout:
a/d/s
- Doom-lite:
w/a/d/space/s
- Flappy:
space
- Pac-Man:
w/a/s/d
- Pong:
w/s
- Overig:
g (terug naar agent), 1-4 (loop count), l (on-device learning toggle), r (restart).
Browser-versie
De no_std int8 engine is ook gecompileerd naar WebAssembly. Deze draait de volledige loop (dream, agent, stream-x learning) in de browser op een canvas.
Reproductie
De volgende stappen kunnen worden gebruikt om het project te reproduceren:
uv sync- Data genereren:
uv run python -m tinyworldformer.training.data --game $G
- Wereldmodel trainen:
uv run python -m tinyworldformer.training.train ...
- Agent trainen:
uv run python -m tinyworldformer.training.streamx --env real ...
- Droom controleren:
uv run python -m tinyworldformer.training.evalwm ...
- Agent in droom trainen:
uv run python -m tinyworldformer.training.streamx --env dream ...
- Exporteren naar Rust:
uv run python -m tinyworldformer.training.export ...
Code-indeling
| Pad | Inhoud |
|---|
src/tinyworldformer/games/ | De vijf games (regels, rendering, read), Game registry, scripted agents |
src/tinyworldformer/training/ | Data, wereldmodel, training, droom-checks, stream-x, export naar Rust |
src/tinyworldformer/rustexport/ | PyTorch → no_std Rust export (QAT layers, int8 spec, codegen) |
tinyworldformer-rs/ | no_std int8 engine, dream loop, agent en learner |
tinyworldformer-rs/firmware/ | Pico 2 W firmware |
tinyworldformer-rs/web/ | WebAssembly build en browser-pagina |
Tiny World Former: Games genereren en spelen op een Raspberry Pi Pico 2 W
Alle frames worden gesampled door het wereldmodel-netwerk, en de acties worden uitgevoerd door het actor-netwerk.
Motivatie
De aanleiding voor dit project was interesse in looped transformers, waarbij één transformer-blok meerdere keren wordt toegepast met dezelfde gewichten in plaats van het stapelen van verschillende lagen. Omdat het hergebruiken van gewichten weinig geheugen kost, was de vraag of een looped world model klein genoeg zou zijn om op een microcontroller te draaien. De keuze viel op de Pico 2 W vanwege de lage prijs en de aanwezigheid van een FPU (Floating Point Unit). Er is gekozen voor eenvoudige games, omdat de echte versies direct in numpy geschreven kunnen worden om elk gegenereerd frame te controleren. Later is een streaming RL-agent toegevoegd, zodat het bord ook kan leren spelen binnen zijn eigen "droom".
Het wereldmodel en de actor
Overzicht van de netwerken
| Kenmerk | Netwerk 1: Het Wereldmodel | Netwerk 2: De Actor |
|---|
| Taak | frame + actie → volgend frame + reward | frame → actie |
| Architectuur | looped transformer, q int8 | MLP met twee verborgen lagen, float32 |
| Grootte | 26k–131k parameters | 9k–16k parameters |
| Kosten per frame | 2–8M multiply-adds | 9k–16k multiply-adds |
| Training | getraind op opgenomen echte games | met stream-x, in het echte spel of in de droom |
De grootte van het model hangt af van het spel; sommige games vereisen meer parameters om speelbaar te worden. Er is gestreefd naar een zo klein mogelijke parameterset.
Op de chip wordt het echte spel door geen van beide netwerken gebruikt: de actor ziet alleen frames van het wereldmodel, en het wereldmodel ontvangt alleen acties van de actor. Beiden draaien op core 1 van de RP2350, terwijl core 0 de USB en het display beheert.
De vooraf getrainde actor kan worden vervangen door een nieuwe stream-x agent die vanaf nul op het bord leert. Bovendien kan de gebruiker via een knopdruk de controle overnemen en zelf in de "droom" spelen.
De spellen
Elk spel bestaat uit enkele honderden regels numpy in src/tinyworldformer/games/. De rendering gebeurt naar een klein raster van token-id's. Eén frame bevat altijd de volledige spelstatus, waardoor de verborgen status in het frame moet worden getekend (bijv. de richting van de bal is onderdeel van het token, en Doom-lite heeft een minimap).
| Spel | Frame | Regels | Random event | Wereldmodel (params / MACs) | Exacte frames |
|---|
| Breakout | 20×16 | paddle, bal, 56 bricks, 5 levens | de serve na een miss | 26k / 4.3M | 99.99% |
| Doom-lite | 32×32 | raycast first-person view van 8×8 doolhof; vijand loopt kortste route; schieten | respawn van de vijand | 55k / 4.2M | 83.6% |
| Flappy Bird | 24×24 | zwaartekracht, flaps, scrollende pipes | hoogte van de volgende opening | 46k / 2.0M | 100% |
| Pong | 20×32 | speler tegen CPU (reageert traag), eerste tot 5 | de serve na een punt | 46k / 2.3M | 99.45% |
| Pac-Man lite | 48×48 | 23×21 doolhof, 249 pellets, twee ghosts op kortste route, 3 levens | respawn van de ghosts | 131k / 8.2M | 86.4% |
"Exacte frames" is het aandeel van achtergehouden overgangen waarbij elke cel van het voorspelde frame exact overeenkomt met het echte spel. Frames direct na een random event zijn uitgesloten, omdat deze alleen gesampled kunnen worden en niet voorspeld. Bij Doom-lite is het percentage lager omdat één foutieve pixel in de 3D-view het hele frame als fout markeert, terwijl de spelstatus bijna altijd correct is.
De echte games kunnen in een terminal worden gespeeld via: uv run python -m tinyworldformer.games.play --game pacman (of breakout, doomlite, flappy, pong).
Hoe het wereldmodel werkt
Frames opgedeeld in tokens
Het frame wordt verdeeld in vierkante patches: 2×2 cellen voor Breakout, 4×4 voor Doom-lite, Flappy Bird en Pong, en 6×6 voor Pac-Man. Dit resulteert in 36 tot 80 tokens per spel. De embedding van een patch is de som van één geleerde vector per cel, met een aparte tabel voor elke positie binnen de patch om de celvolgorde te behouden. Elke token krijgt daarnaast een positievector en de vector van de actie.
Looped Transformer
Het model is een enkel transformer-blok (attention plus een ReLU MLP) met één set gewichten, toegepast 1 tot 4 keer, waarbij de input voor elke pass opnieuw wordt toegevoegd. De output na elke pass is getraind tegen het volgende frame. Tijdens runtime is het aantal loops een afweging tussen snelheid en nauwkeurigheid. Op de Pico worden twee loops gebruikt om een acceptabele framerate te behouden.
Exacte achtergehouden frames per aantal loops:
| Loops | Breakout | Doom-lite | Flappy Bird | Pong | Pac-Man |
|---|
| 1 | 99.68% | 59.8% | 99.99% | 98.47% | 79.6% |
| 2 | 99.99% | 83.6% | 100% | 99.45% | 86.4% |
| 3 | 100% | 85.9% | 100% | 99.60% | 85.7% |
| 4 | 99.96% | 84.7% | 100% | 99.53% | 84.4% |
Decoding
Het model produceert een score voor elke cel en elk token. Het omzetten van deze scores naar een frame gebeurt via drie regels:
- Achtergrondcellen (muren, bricks, pellets, vloer): elke cel neemt onafhankelijk het meest waarschijnlijke token.
- Singletons (de bal, de speler, elke ghost): deze worden geplaatst via één keuze over alle (cel, token) paren, zodat er altijd precies één van elk is. Na een random event wordt de keuze gesampled tussen kandidaten binnen 2.2 logits van de beste optie.
- Sprites: Pac-Man en de ghosts zijn 2×2 sprites, maar slechts één cel draagt de status (bijv. richting). De overige drie cellen worden hiervan afgeleid in plaats van voorspeld.
Validatie van het wereldmodel
Om te controleren of het model robuust is, speelt een agent in de droom. De spelstatus wordt teruggelezen uit het gedroomde frame, de echte regels worden toegepast op die status en actie, en het resultaat wordt vergeleken met het volgende gedroomde frame.
Resultaten over 5.000 gedroomde stappen (bij 2 loops):
| Spel | Foutieve spelstatussen | Foutieve rewards | Foutieve cellen per stap |
|---|
| Breakout | niet teruggelezen; 0 frames met onmogelijke tokens | 6 | – |
| Doom-lite | 1 | 0 | 1.50 |
| Flappy Bird | 0 | 0 | 0.00 |
| Pong | 1 | 0 | 0.00 |
| Pac-Man | 742 | 38 | 0.39 |
Training
- Dataverzameling: Agents spelen het echte spel; elke stap wordt opgenomen als een voorbeeld (frame, actie, volgend frame, reward). Er worden 1 tot 3 miljoen voorbeelden per spel verzameld.
- Training Wereldmodel: Het model leert het volgende frame te voorspellen met een per-cel cross-entropy loss (30k stappen). De eerste 3–4k stappen zijn in float, daarna schakelt Quantization Aware Training (QAT) in voor int8 rounding.
- Implementatie: De int8-gewichten worden in de Rust-broncode geschreven. De firmware draait vervolgens op de Pico.
- Fine-tuning: Een agent die in het wereldmodel is getraind, kan statussen bereiken die niet in de oorspronkelijke data zaten. Er worden nieuwe echte games opgenomen met deze agent, waarna het wereldmodel daarop wordt gefinetuned.
Gevonden exploits en fixes
- Dubbele objecten: Het wegen van de loss op basis van "cel verandering" zorgde ervoor dat het model twee ballen tekende. Dit is opgelost door weging op basis van waar actie mogelijk is in het inputframe en de globale singleton-keuze.
- Logische fouten: Een Breakout-agent scoorde 154 punten terwijl er 26 bricks overbleven, omdat het model bricks in de hoeken niet verwijderde bij impact.
- Droom vs. Realiteit: Agents die alleen in de droom trainden, scoorden daar 210 (max realiteit is 56) maar slechts 28 in het echte spel.
- Sprite-verlies: In Doom-lite verdween de vijand soms uit de droom. Noise-augmentatie (random pixels toevoegen aan input views) heeft dit opgelost.
De agents: Stream-X
De agents maken gebruik van Stream AC($\lambda$), wat leert van één transitie per keer zonder replay buffer, batches of target networks. Eligibility traces sturen credit terug in de tijd, en ObGD voorkomt dat stappen het TD-target overschrijden. De policy- en value-netwerken hebben elk twee verborgen lagen van 64 units (32 op de Pico) met LayerNorm.
Input features per spel:
- Breakout: rij, kolom en richting van de bal.
- Pac-Man: voor elke richting de kortste afstand tot de dichtstbijzijnde pellet en de afstand tot elke ghost.
Greedy return in het echte spel (gemiddelde van 10 episodes):
| Spel | 0 | 10k | 50k | 100k | 250k | Max |
|---|
| Breakout | 8.2 | 29.9 | 29.9 | 31.1 | 45.6 | 56 |
| Doom-lite | −5.0 | −5.0 | 69.9 | 78.1 | 82.9 | – |
| Flappy Bird | −1.0 | 24.6 | 96.2 | 123.7 | 332 | 332 |
| Pac-Man | 7.0 | 17.7 | 72.9 | 146.6 | 217.4 | 249 |
| Pong | −4.3 | −2.0 | −0.3 | 1.1 | 1.1 | 5 |
Om te voorkomen dat Doom-lite agents alleen in de hoek blijven staan schieten, is een --idle-penalty 0.1 toegevoegd.
Quantisatie en de Pico
Elke matrix in het wereldmodel is int8 met één schaal per output-rij. Activaties zijn eveneens int8. Matrixvermenigvuldigingen sommeren int8 × int8 naar int32 (gebruikmakend van de SMLAD instructie op de RP2350). Softmax maakt gebruik van een exp-tabel van 1024 entries. Norms, residuals en embeddings blijven in float32 vanwege de FPU op de RP2350.
Geheugengebruik
| Spel | int8 gewichten | float32 (emb, norms, scales) | Agent (float32) | Firmware Flash | Firmware RAM |
|---|
| Breakout | 19.7 KB | 31.6 KB | 54.1 KB | 361 KB | 299 KB |
| Doom-lite | 33.9 KB | 91.9 KB | 61.3 KB | 471 KB | 357 KB |
| Flappy Bird | 30.1 KB | 69.5 KB | 50.8 KB | 409 KB | 227 KB |
| Pong | 30.1 KB | 70.4 KB | 52.6 KB | 414 KB | 244 KB |
| Pac-Man | 131k params | - | - | 600 KB | 460 KB |
Op de chip
De firmware is geschreven in no_std Rust op Embassy. Er vindt geen runtime-allocatie plaats; gewichten staan in flash en buffers zijn vaste statische variabelen in RAM.
- Core 1: Draait de game loop. De actor kiest een actie → het int8 transformer model voert 2 loops uit → decoding zet scores om in een frame → de reward wordt bepaald.
- Core 0: Beheert async taken. Verstuurt frames via USB-serieel als gekleurde ANSI-tekst en beheert het optionele Display Pack (ST7789 over SPI).
Geheugenverdeling:
- Flash: Wereldmodel (int8 matrices, float32 embeddings), actor (MLP gewichten) en firmware code.
- RAM: Scratch buffers voor forward passes, features en hidden layers van de actor, en weights/traces voor de on-device learner.
Flashing en Uitvoering
Het flashen gebeurt via het tinyworldformer-rs crate. Voorbeelden van commando's:
./flash.sh --env doomlite --agent doomagentroam_p0.3 (bouwen en flashen)./flash.sh --env breakout --display (met Display Pack)./flash.sh --env pong --host (spelen in de terminal)
Besturing via USB:
- Breakout:
a/d/s
- Doom-lite:
w/a/d/space/s
- Flappy:
space
- Pac-Man:
w/a/s/d
- Pong:
w/s
- Overig:
g (terug naar agent), 1-4 (loop count), l (on-device learning toggle), r (restart).
Browser-versie
De no_std int8 engine is ook gecompileerd naar WebAssembly. Deze draait de volledige loop (dream, agent, stream-x learning) in de browser op een canvas.
Reproductie
De volgende stappen kunnen worden gebruikt om het project te reproduceren:
uv sync- Data genereren:
uv run python -m tinyworldformer.training.data --game $G
- Wereldmodel trainen:
uv run python -m tinyworldformer.training.train ...
- Agent trainen:
uv run python -m tinyworldformer.training.streamx --env real ...
- Droom controleren:
uv run python -m tinyworldformer.training.evalwm ...
- Agent in droom trainen:
uv run python -m tinyworldformer.training.streamx --env dream ...
- Exporteren naar Rust:
uv run python -m tinyworldformer.training.export ...
Code-indeling
| Pad | Inhoud |
|---|
src/tinyworldformer/games/ | De vijf games (regels, rendering, read), Game registry, scripted agents |
src/tinyworldformer/training/ | Data, wereldmodel, training, droom-checks, stream-x, export naar Rust |
src/tinyworldformer/rustexport/ | PyTorch → no_std Rust export (QAT layers, int8 spec, codegen) |
tinyworldformer-rs/ | no_std int8 engine, dream loop, agent en learner |
tinyworldformer-rs/firmware/ | Pico 2 W firmware |
tinyworldformer-rs/web/ | WebAssembly build en browser-pagina |
Groetjes,