# Replacing my game engine with a neural net

I make complex ASCII games and I wondered, since each frame is a grid of characters, *could a ML model predict the next frame*? What would it be like to play the game inside of the model's "imagination?" Would it even be possible to train such a model from scratch?

I've never built my own model before, so I embarked on the journey with Claude by my side. The plan was to build a model locally on my laptop, taking inspiration from other world models like Google's [Genie](https://deepmind.google/models/genie/) and World Lab's [Atlas](https://worldlabs.ai/blog/atlas), and other game frame prediction experiments like Google's [GameNGen](https://gamengen.github.io/) and NVIDIA's [GameGAN](https://research.nvidia.com/labs/toronto-ai/GameGAN/). We experimented with CNNs, RNNs, diffusion models, and transformers.

TL;DR: it worked! We built a 0.6M-parameter CNN trained on 200k frames that plays the game in the browser in real time. Yes there are caveats – new terrain loses details, continuity is nearly non-existent, and NPC behavior is quite limited. But the concept works impressively well given how small the model and training data is.

> [You can play the game in the neural net yourself here](https://thunder-lizard-frame-prediction.vercel.app/) (requires a browser with WebGPU like Chrome or Edge)

Every frame you see is drawn by the model. There is no engine running behind it.

[![Playing Thunder Lizard inside the neural net](https://cdn.hashnode.com/uploads/covers/64c6cae04e2054e7a79704e7/8200614a-0ddb-4ea7-a367-b78034fc102b.gif align="center")](https://thunder-lizard-frame-prediction.vercel.app/)

### But why?

Why even bother with this? What is the benefit? Would I really replace the engine with a model? Definitely not. The experiment was more conceptual, to see what is possible, and to learn in a very hands on way what it is like to build my own model, and all the tools and techniques that goes into it. This also was a nice taste of the bleeding edge of experimentation that the leading labs of the world are doing at a way larger scale right now. And it gives some insights into other applications with similar techniques, like weather modeling. And, of course, because it is fun!

---

## Deep Dive

The rest of this post dives into the deeper details of how it came together and what I learned. All the code and experiments are available in the [github repo](https://github.com/jschomay/thunder-lizard-next-frame-prediction). This deep dive only scratches the surface of the full experiment, so if you really want the nitty gritty details, check out the full [experiment journal](https://github.com/jschomay/thunder-lizard-next-frame-prediction/blob/main/JOURNAL.md).

### The game

The original game is something I made a while ago. You are a dinosaur on a procedurally generated prehistoric island. You eat smaller dinos, run from bigger ones, and stay ahead of the lava that slowly spreads across the island. Play the original at: [Thunder Lizard: Eat or Be Eaten](https://enegames.itch.io/thunder-lizard-eat-or-be-eaten).

### Generating data

A model learns from examples, so step one was getting a lot of them, much more than what I could play manually. Claude made a headless version of the game with no rendering output for bots to play. Each tick recorded the grid of characters and any key pressed, which turns into about 2,000 cells per frame, for about 200,000 frames.

### Representing the data

At first I thought we could just feed the characters to a LLM, but that wouldn't work because the grid of characters has no semantic meaning to the LLM so it wouldn't know what to do with it. Instead we went with a CNN, which required representing all the data numerically. The first step was classifying each cell into its terrain type with special classifiers for dinos and their positions.

![Representing a frame as numbers](https://cdn.hashnode.com/uploads/covers/64c6cae04e2054e7a79704e7/df6e0716-8c58-4558-ac19-ac0ec728b54a.png align="center")

I learned that plugging raw numbers into the model doesn't work. Instead, each number is represented as a "one-hot" vector which means each cell has a vector of size 8 for the 8 terrain types, with only one of them set to 1 and the others set to 0. It also gets a similar one-hot vector for any dinos in the space. Plus we included the same info for previous frames to show changes over time. We experimented with how many previous frames to include. The image below represents data for 1 cell with 2 frames of history.

![A single cell stores current and previous terrain and entity data as one-hot vectors](https://cdn.hashnode.com/uploads/covers/64c6cae04e2054e7a79704e7/348fc501-75f7-442d-828a-911f141434f8.png align="center")

### Training the model

We started with a convolutional neural network (CNN), chosen because it works well for images or grids of independent cells. It looks at a patch of cells around each cell and decides what that cell becomes next. We gave it all those one-hot vectors for each cell in the frame with history and scored the predicted frames against actual frames for batches of our gathered data. One key early insight was that most cells stay the same from frame to frame, so we focused on the cells that changed.

A single run took about 20 minutes on my M3 Macbook Air. We recorded everything in Weights & Biases.

![training data graphs in Weights & Biases](https://cdn.hashnode.com/uploads/covers/64c6cae04e2054e7a79704e7/d941d353-d63e-4d8f-85a5-7bcb8671325e.png align="center")

### Validation

Claude made a number of programmatic scorers for different kinds of accuracy, but the real test was putting the model into a playable harness. We did this by exporting the model to [ONNX Runtime Web](https://onnxruntime.ai/) on WebGPU. Each frame of gameplay the harness does the following process:

1. Encodes the current and past frames and any player input into the representation described above
2. Passes it through the model to get the predicted output
3. Decodes it back into a frame representation
4. Renders the frame as the matching character glyphs of the original game
5. The new frame becomes the input for the next tick and the process repeats

Our model was small enough that all of this happened in less than 16ms. This let me play the game in the neural net in realtime. The first time I played was magical. The game looked like what I was used to, but moving around was more dream-like, as terrain shifted around, mostly resembling what it should, but with some "cracks" in reality.

### Learnings and further experiments

The process was very iterative, identifying issues, forming hypothesis and validating or invalidating them. For example, dinos originally exponentially multiplied until the screen was full of them! We determined it was because each cell independently determines what it should show, so multiple cells were deciding a nearby dino should move to their spot. We solved this treating the dinos separately than terrain, predicting the next move each dino would take rather than which cell would be its new location. The [hypothesis tracker](https://github.com/jschomay/thunder-lizard-next-frame-prediction/blob/main/HYPOTHESES.md) shows more examples.

The two main issues we encountered were maintaining continuity if you move away from an area and return, and generating realistic terrain the further you walk. Because cells determine their type from their surroundings, they tend to stabilize on large patches of the same terrain over time. And because we have such a short history given to the model, consistency isn't a concept it learns.

We went through a number of approaches to attempt to improve on these shortcomings. Simpler approaches were to create more accurate training data in larger volumes or to train for longer. We moved the training to Colab to speed that part up. We also experimented with other model architectures.

We used a U-Net to give cells a larger viewport of the world around them by adding more layers to the model, but it didn't make a big difference. We used a RNN to add recurrent memory to help with consistency, but because of model size limits it wasn't enough. We tried diffusion models to get better generation, but instead it drifted towards blotches of common terrain. We used a transformer model with a much longer history, which had better scores on paper, but in actual play was not a big improvement considering the dramatic model size increase.

We tracked all of these parameters in a parallel coordinates graph to find the right mix that had the best properties, but ultimately everything became a tradeoff in one direction or another.

![Parallel coordinates of runs vs scores](https://raw.githubusercontent.com/jschomay/thunder-lizard-next-frame-prediction/main/media/wandb-parallel-coordinates.png align="center")

### Conclusions

I would call this experiment a success. I was able to play the game in the model and see what it was like. For a moment, I was a tiny speck moving in parallel with the giants of the ML world, who face the very same problems I encountered. My 200,000 frames and 600,000 parameters held up surprisingly well when compared to GameNGen's 900 million frames and 860 million parameters.

Most importantly, I learned so much about ML that previously had only been abstract concepts I read about. I feel much more empowered to continue building custom models for other pursuits.