Teaching a World Model to Play Pokemon
A developer trained a small 12.5-million parameter world model locally using a JEPA architecture to successfully plan button sequences that choose a starter in Pokémon Red.
read it Offers a clear, hands-on technical breakdown of training a local JEPA world model with SIGReg regularization and rollout fine-tuning.
summary
The author built a 12.5-million parameter Joint Embedding Predictive Architecture (JEPA) world model to learn the state dynamics of Pokémon Red without using rewards. By using SIGReg regularization to prevent latent collapse and fine-tuning on self-generated rollouts, the model successfully planned a 14-button sequence to select a starter Pokémon. Testing 100 fresh seeds showed the model achieved a 52% success rate, outperforming random button presses and un-fine-tuned models.
key points
- Trained a 12.5M parameter LeWorldModel1 variant locally on an RTX 3080 Ti using 42,382 grayscale frames from Pokémon Red recordings.
- Prevented latent collapse using SIGReg (Sketched Isotropic Gaussian Regularization) with characteristic function tests over random directions.
- Addressed error compounding during planning by implementing rollout fine-tuning, allowing the predictor to learn from its own generated states.
who should read it
Machine learning engineers and hobbyists interested in world models, JEPA architectures, and offline model training.
discussion positive
consensus
Commenters find the project an engaging practical exploration of world models, discussing how it compares to traditional reinforcement learning and whether the architecture can scale to full game completion.
pushback
A few commenters questioned the practical necessity, noting that frontier models already possess general knowledge to solve Pokémon without custom training.
I think it is possible but seems quite difficult to me from this experiment. The hard part is that Pokemon has a lot of different goals, and has many strategies to complete it. You would need a way to encode these intermediary goals in some way temporally that the model needs to learn to complete in order. For example, you need to collect all eight gym badges before you can challenge the Elite Four to win.
You could compose multiple network, one with the goal of learning how to interact with the world through the screenshots. This would probably require a lot of frames of the game from many …
The author expands on the core scaling challenge, explaining why managing multi-step intermediate goals like gym badges makes full game completion difficult for this setup.
Generated by an LLM from the article and thread. Check the sources before relying on it.