deniz.in

Markets

Weather

Loading weather

· via Hacker News – Front Page (hnrss.org)

JEPA world model trained on one GPU plans its way to a starter in Pokémon Red

A nostalgia.dev post details training a small JEPA-style world model on a single RTX 3080 Ti to predict Pokémon Red's dynamics, with the final model planning the button presses that select Squirtle.

JEPA world model trained on one GPU plans its way to a starter in Pokémon Red

A project posted on nostalgia.dev and surfaced on the Hacker News front page documents a small but instructive experiment in model-based planning: teaching a world model to play the opening of Pokémon Red, culminating in the model planning a sequence of button presses that selects Squirtle as its starter Pokémon.

A deliberately narrow task

According to the post, the original goal was to reach Professor Oak's lab, finish his dialogue, pick a starter, leave the lab and defeat Oak's grandson. The author quickly judged that too ambitious for a first experiment and narrowed the objective: starting from a saved state inside Oak's lab, successfully acquiring any of the three starters — Bulbasaur, Charmander or Squirtle — would count as success.

Even that reduced task is not trivial. From the saved position, pressing the A button twelve times in a row is enough to complete the selection. But an agent with no understanding of the game could just as easily wander the lab indefinitely, or press B partway through and cancel its own progress.

What the model actually learns

The author defines a world model as a function that maps a current observation and an action to the next observation, written schematically as the next state approximating a function of the current state and the button pressed. The point is not to memorise that pressing left works in one specific room, but to generalise the intuition that pressing left moves the player left anywhere in the game.

The setup resembles reinforcement learning with one crucial difference the post highlights: the world model learns state dynamics without any reward signal, making it reward-free.

Rather than predicting raw screenshots, the model works in latent space. An encoder converts each screenshot into a 192-number embedding — a learned vector representation whose individual coordinates carry no human-assigned meaning. A separate predictor takes the current embedding plus the chosen action and outputs a predicted next embedding, which is compared against the embedding of the screenshot that actually followed using mean squared error. The predicted embedding is never decoded back into an image; everything happens in the learned representation.

The collapse problem

Training an encoder and predictor together hides what the post calls a dangerous degenerate case. Because the same encoder produces both the input to the predictor and the ground-truth target, the system can satisfy the prediction loss in a way that strips the embedding of useful information. The post addresses this with SIGReg, a regularisation technique it credits with keeping the embedding space useful — the classic representation-collapse problem that latent-prediction architectures such as JEPA must solve to be trainable at all.

The architecture itself is a variant of Yann LeCun's JEPA (Joint Embedding Predictive Architecture), specifically one the author found appealing because it was small enough to train locally on a single RTX 3080 Ti and simpler in design than many earlier JEPA models.

Setbacks before success

The experiment did not go smoothly. The author reports more difficulties and setbacks than expected, even though the model had looked promising in simpler tests. The post's structure indicates the first planning attempt failed, prompting a round of rollout fine-tuning before a second attempt — training not just on single-step predictions but on multi-step sequences the model generates itself. The final result, shown in the post, is the trained model planning the button presses that select Squirtle from the table in Oak's lab.

Why it matters

World models are one of the more closely watched directions in AI research, and this post is a useful ground-level view of why. It demonstrates that a JEPA-style model small enough for a consumer GPU can learn usable environment dynamics from nothing but screenshots and button inputs, and then be queried for a plan — no reward function, no hand-coded game logic.

Just as valuable is the honest accounting of failure modes: latent prediction invites representational collapse, and a model that predicts one step accurately can still fall apart over the longer rollouts that planning requires. Those are the same tensions the field is working through at scale, reproduced here in a 1996 Game Boy world — which makes the post a readable bridge between research papers and practice for anyone tracking model-based agents.

  • #world-models
  • #jepa
  • #machine-learning
  • #game-ai
  • #pokemon