Teaching a World Model to Play Pokemon
6 hours ago
- The author trained a JEPA-based world model (LeWorldModel) on Pokémon Red screenshots to plan button sequences for selecting a starter Pokémon.
- World models learn state-transition dynamics from observations without rewards; the model predicts future embeddings rather than raw screenshots.
- To prevent latent collapse (all embeddings becoming identical), SIGReg regularization encourages embeddings to resemble an isotropic Gaussian.
- Training data consisted of 42,382 grayscale frames grouped into 1,009 short trajectories, including both clean and noisy routes.
- Planning uses predicted embeddings from a 14-step rollout and the cross-entropy method (CEM) to search for sequences that minimize distance to goal embeddings.
- Initial plans failed due to compounding prediction errors; rollout fine-tuning (training the predictor on its own predictions) improved performance.
- After fine-tuning, the model achieved 52% success rate in selecting a starter (Squirtle) from a saved state in Oak's Lab, compared to 0% for random actions.
- The final model had ~12.5 million parameters, trained end-to-end from scratch, and remained limited to short tasks due to exponential difficulty scaling.