I think some of the most fascinating work happening in the field of Artificial Intelligence surrounds world models. There are many kinds of world models (and the term itself has become a bit overloaded), but one of the most exciting architectures for world models is Yann LeCun’s JEPA (Joint Embedding Predictive Architecture). There are many variants, and one that caught my eye in particular was LeWorldModel. It seemed small enough to train locally on my RTX 3080 Ti and had a simpler design than many of the previous JEPA architectures.
San Francisco has been covered in Pokémon memorabilia–big advertisements featuring many of the 1026 pocket monsters plastered across the subway lines and bus stops on Market Street. SF was the site of the Pokémon World Championships this year, and many eclectic and joyous Pokémon trainers could be found wandering the temperate hills and concrete financial district of downtown SF. Maybe all the advertisements had subliminally controlled me, but I had decided that a good “world” for our world model would be the 1996 game that started it all, Pokémon Red.
The mechanics of Pokémon
Pokémon Red is a game where you explore a world, collect creatures called Pokémon, and use them in battles.
You use the direction buttons to walk around or move through menus. The A button interacts with things, advances dialogue, and confirms choices; the B button usually cancels or backs out. Choosing a starter means approaching a Poké Ball (a container that holds Pokémon) and getting through the dialogue that confirms your choice.
Near the beginning of the game, the player is inside Professor Oak’s lab, where he offers you your first Pokémon: Bulbasaur, Charmander, or Squirtle. These are the three “starters,” each waiting in a Poké Ball on a table in his lab.
The GIF above is the final result of the model training: the model planned a sequence of button presses that selected Squirtle. There were more difficulties and setbacks than I expected, even though the model had looked promising in simpler tests.
The Plan
The goal was relatively straightforward: defeat Professor Oak’s grandson. Breaking it down into multiple steps, I came up with:
- Get to Professor Oak’s Lab
- Finish Dialogue with Oak
- Select a Starter Pokémon
- Try to exit the lab
- Face and beat Oak’s grandson
It quickly came to my attention that this might be ambitious for a first experiment–so I narrowed it down to selecting a starter from a saved state in Oak’s Lab, where acquiring any of the three starters would count as a successful attempt.
From the saved position in the lab, having the model press A twelve times is enough. But can the model learn that? The model could also just wander around aimlessly, maybe in perpetual torment inside Oak’s lab. Or the model could press B after every few sequences of A, cancelling its effort when it almost reached its goal.
What Even is a World Model?
A world model is a model in which we begin with some current state or observation, and some action happens that modifies this state, producing a new observation. Realistically, there should be some sort of correlation or hopefully causation between the action on some state and the new state it produces. The goal of the world model is to learn this correlation. Let’s say the current observation is a screenshot of Professor Oak’s lab, with the player in front of the Poké Balls, and the action is pressing left, moving the player one tile to the left–producing the new end screenshot. The goal would be for the world model to develop some form of intuition that pressing left moves the player left.
$$ o_{t+1}\approx F(o_t,a_t). $$Here \(o_t\) is the current screenshot, \(a_t\) is the button pressed, and \(F\) is the function the world model learns to predict the next screenshot \(o_{t+1}\).
But the goal should be to generalize–the model shouldn’t learn that pressing left in Professor Oak’s lab moves the player left, but that pressing left anywhere should move the player to the left.
As an aside, this setup might sound somewhat similar to reinforcement learning. But crucially, the world model learns to understand the state dynamics without a reward, aka being reward-free (more on this later).
How the World Model Learns to Predict
From screenshots to embeddings
So if what we are really interested in is the state-transitions and the dynamics of the environment and all we have are observations (in this case, in the form of screenshots from our game), does the model learn to predict screenshots?
No, what the model actually learns is to predict within its latent space, also known as its embedding space. We first need to take an encoder that, when given a screenshot, produces an embedding.
What is an embedding?
An embedding is a learned representation of an input as a vector of numbers. An encoder can turn an image, a sentence, or a sound into such a vector, giving another model something it can compare, predict, or use as input.
The entries are not hand-labeled features: one coordinate does not have to mean “color” or “position.” Information can be spread across many entries, and some details of the original input can disappear altogether. Whether an embedding is useful depends on which distinctions the encoder learns to keep.
With the embedding representing the current observation \(o_t\), we can then have a predictor model that, given an action \(a_t\), tries to predict the future embedding. With this future predicted embedding, we can try to compare it to the real embedding produced by the future observation \(o_{t+1}\). This is actually rather simple: just pass the new screenshot through the same encoder and we get the real future embedding. The goal is then to just have:
$$ \hat z_{t+1}\approx z_{t+1}. $$Here, \(\hat z_{t+1}\) is the predicted next embedding, and \(z_{t+1}\) is the embedding of the screenshot that actually followed.
Writing this out for one step:
$$ z_t = E(x_t), \qquad \hat z_{t+1} = P(z_t,a_t), \qquad z_{t+1} = E(x_{t+1}). $$Here, \(x_t\) is the current screenshot, \(a_t\) is the action, and \(E\) and \(P\) are the encoder and predictor. Importantly, the predicted embedding \(\hat z_{t+1}\) is never turned into a screenshot, and we can compare it directly with \(z_{t+1}\).
Prediction loss and collapse
To learn the difference between our prediction and the ground-truth embedding, we will use the ol’ reliable loss function Mean-Squared Error.
$$ \mathcal L_{\mathrm{pred}}=\frac{1}{D}\sum_{d=1}^{D}(\hat z_{t+1,d}-z_{t+1,d})^2,\qquad D=192. $$It looks kind of complicated but you can imagine it as just a distance function. We are basically just trying to measure the distance between our two embedding vectors.
Everything seems simple enough: minimize this function–we have a world model that can play Pokémon, right?
But it’s never so simple, and there is actually a dangerous degenerate case that can happen. Remember that we are training the encoder and the predictor, and the encoder appears twice in each training example:
- The encoder turns the current screenshot into \(z_t=E(x_t)\).
- The predictor takes \(z_t\) and the button \(a_t\), then guesses \(\hat z_{t+1}=P(z_t,a_t)\).
- The same encoder turns the actual next screenshot into \(z_{t+1}=E(x_{t+1})\).
- Training reduces the distance between \(\hat z_{t+1}\) and \(z_{t+1}\), updating both the predictor and the encoder.
In the worst case, the encoder starts to learn to embed every screenshot as the same embedding \(c\). Then, the predictor also begins to learn to predict that the next embedding will be \(c\). This is called latent collapse and means that the model has found a trivial way to reduce the loss function–by collapsing all embeddings to the same embedding.
SIGReg: keeping the embedding space useful
So the prediction loss can be minimized by putting every screenshot into the same point. If we could somehow keep the embeddings from collapsing to the same point, then we can get back on track at producing a useful world model. This is the perpetual problem of JEPA, and many different JEPA model architectures have come up with different ways to try to prevent latent collapse.
One possible fix is to make sure each number in an embedding changes across screenshots. But we have to be careful about how the embeddings spread out. Imagine embeddings with just two numbers: \((1,1)\), \((2,2)\), and \((3,3)\). Both numbers change, but the second always copies the first. The points form a line, which only means it has one dimension of data instead of spreading over a two-dimensional area. With 192 numbers, the same kind of redundancy is harder to notice.
The LeWorldModel paper chooses a relatively simple solution–at least compared to many of the other JEPA architectures–called SIGReg, or Sketched Isotropic Gaussian Regularization.
SIGReg asks for a stronger shape. Across a batch of screenshots, it encourages the embeddings to resemble a standard isotropic Gaussian:
$$ z\sim\mathcal N(\mu,\Sigma),\qquad \mu=\mathbf 0,\quad \Sigma=I_D. $$The isotropic Gaussian has a mean of 0, and the identity matrix \(I_D\) as its covariance matrix. In two dimensions, that is a round cloud rather than a point or a line. You can see here that the isotropic Gaussian forms this more evenly spread out shape in all directions:
Suppose the encoder produces a batch of embeddings \(z_1,\ldots,z_N\in\mathbb R^D\) which form a cloud like in the above figure. SIGReg picks a random direction \(u\) of length one and measures where each embedding falls along it:
$$ h_i=u^\top z_i,\qquad i=1,\ldots,N. $$The \(h_i\) are now ordinary numbers instead of \(D\)-dimensional vectors. If the original cloud really followed \(\mathcal N(0,I_D)\), the numbers along any unit direction would follow \(\mathcal N(0,1)\).
Why does every direction give a standard Gaussian?
In two dimensions, let a random embedding be \(X=(X_1,X_2)\), where \(X_1\) and \(X_2\) are independent standard Gaussians. Choose the unit direction \(u=(3/5,4/5)\). Its projection is
$$ h=u^\top X=\frac35 X_1+\frac45 X_2. $$A weighted sum of independent Gaussians is still Gaussian, and its mean is zero. To find its variance, remember that variance averages squared deviations from the mean. Multiplying \(X_1\) by \(3/5\) therefore multiplies its variance by \((3/5)^2\); the same applies to \(X_2\). The mixed term averages to zero because the coordinates are independent and centered at zero, so their variance contributions add:
$$ \operatorname{Var}(h)=\left(\frac35\right)^2+\left(\frac45\right)^2 =\frac9{25}+\frac{16}{25}=1. $$The same argument works in \(D\) dimensions. For any unit direction \(u=(u_1,\ldots,u_D)\), its squared components add to one: \(\sum_j u_j^2=1\). The projection \(u^\top X=\sum_j u_jX_j\) is Gaussian with mean zero and variance \(\sum_j u_j^2=1\), so it follows \(\mathcal N(0,1)\). In compact matrix notation, that variance calculation is \(u^\top I_Du=\|u\|^2=1\).
The problem has become a one-dimensional question: do these projected numbers look like samples from a standard Gaussian?
SIGReg answers this with a characteristic-function test.
The Characteristic Function
A characteristic function is another way to describe the shape of a probability distribution.
For a random vector \(X\), it is defined as
$$ \phi_X(t)=\mathbb{E}[e^{i t^\top X}]. $$You can think of \(t\) as choosing a direction and a scale at which to “probe” the distribution. The characteristic function tells us what the distribution looks like under that probe.
For an isotropic Gaussian,
$$ X\sim\mathcal N(0,\sigma^2 I), $$the characteristic function is
$$ \phi_X(t)=e^{-\frac12\sigma^2\|t\|^2}. $$Notice that this only depends on \(\|t\|\), the length of \(t\), and not the direction it points. That matches the intuition behind an isotropic Gaussian: it looks the same in every direction.
The value of the characteristic function is that it is unique to each probability distribution, and therefore if you can show that a random variable’s characteristic function matches another random variable’s everywhere, then they have the same distribution.
For each direction, SIGReg compares the characteristic function estimated from the projected samples with the known one for \(\mathcal N(0,1)\). We then use the Epps–Pulley test to turn the mismatch into a penalty. You can find more details in the original LeJEPA paper or Appendix A of LeWorldModel.
We average that penalty over 1,024 random directions, then use its gradient to update the encoder. The finite set of directions is then used as an approximation to checking the full distribution.
The final training objective function looks like this:
$$ \mathcal L=\mathcal L_{\mathrm{pred}}+0.1\,\operatorname{SIGReg}(Z). $$Here, the 0.1 controls how much the regularizer contributes to our training by discouraging the erasure of variations among screenshots.
The Pokémon Training Data
We now have a way to train a model from pairs of screenshots and actions. Where do those pairs come from? I recorded 42,382 grayscale frames from the Pokémon Red emulator, grouped into 1,009 short trajectories. Some trajectories follow scripted routes to a starter. Others add noisy actions to those routes, or move around more randomly. For each step, the recording contains the current frame, the button pressed, and the frame that followed.
It’s actually important to have the messy trajectories. A model trained only on the clean route might see A pressed whenever a dialogue box appears and never learn what B does there. The planner, however, is going to propose all sorts of button sequences, including bad ones, and it needs predictions for those sequences too. This does not make the dataset a complete map of Pokémon Red, but it gives the model more than a single demonstration to memorize.
Notice that none of these recordings tells the world model whether a trajectory succeeded. The training loss asks it to predict what follows a button press; it never pays the model for acquiring a Pokémon. The goal will enter later, when we use the trained model to plan.
Before planning, I wanted to know whether the embeddings kept track of whether the player had chosen a Pokémon. I froze the encoder and trained a small classifier to answer that question from the embeddings, using the emulator’s party count as the correct answer. This is a linear probe: the encoder cannot change, so the classifier has to work with information the model already kept.
On gameplay recordings that weren’t used to train the classifier, it could still tell whether the player had a Pokémon. I also checked the dynamics: the predictor did better than a baseline that simply copied the current embedding, and giving it the wrong button made its prediction worse. These checks were promising, but I still needed to find out whether the model could plan a whole sequence.
Planning in the Learned World
To ask for a starter, I took screenshots from successful Bulbasaur, Charmander, and Squirtle selections and passed each through the encoder. Their embeddings form a set of three possible goals, \(\mathcal G\). Any one of them will do. This gives the planner a picture of a desired outcome without supplying a route through the lab or a reward for each button press.
Now imagine trying a sequence of buttons without actually pressing them. We start with the real screenshot from the saved state, encode it once, and give that embedding and the first proposed button to the predictor. For the second button, there is no new screenshot: we have to feed the first prediction back in. In the simplified notation from earlier, a candidate plan unfolds as
$$ \hat z_0=E(x_0),\qquad \hat z_{t+1}=P(\hat z_t,a_t). $$The actual predictor can use a short history of states, but the important feature is the same: after \(\hat z_0\), the imagined states come from the model itself. We can roll out many proposed action sequences this way without running the emulator for every one.
We need one score for each imagined fourteen-button sequence. Here is how we calculate it from the predicted states and the three starter goals:
$$ J(a_{0:H-1})=\min_{1\leq k\leq H}\;\min_{z_g\in\mathcal G}\frac{1}{D}\|\hat z_k-z_g\|_2^2,\qquad H=14. $$It looks complicated, but we can walk through it one button at a time. After each imagined button press, we compare the predicted embedding with the three embeddings from successful starter selections and keep the distance to whichever starter is closest.
Once we have done that for all fourteen steps, we keep the smallest distance we saw. That number is \(J\). A plan can score well by getting close to a starter at step 12, even if it continues for two more actions. Of course, this is only what the model predicts: to find out whether the player actually got a Pokémon we have to run the predicted plan directly in the emulator.
There are too many fourteen-button sequences to try one by one, so I used the cross-entropy method, or CEM, to narrow the search. Picture a plan as fourteen empty slots. For each slot, CEM keeps track of how likely it is to put each possible button there. At first, it samples a wide variety of plans.
Each round works like this:
- Sample 512 complete plans and use the world model to predict what each one would do.
- For each plan, compare its fourteen predicted steps with the three starter goals. The smallest embedding distance is its \(J\) score. Keep the 64 plans with the lowest scores.
- Look at those 64 plans, slot by slot, and make their buttons more likely to be sampled in the next round.
For example, if many of the 64 plans press A as their fourth action, A becomes more likely in the fourth slot next round.
We repeat this process for several rounds, gradually favoring plans with lower predicted costs. Then we run the best plan in the emulator.
Why the First Plan Failed
The first search found a plan that the model thought would get close to a starter, but in the emulator, the same buttons left the player without a Pokémon. Somewhere between the imagined sequence and the real one, the model had gone wrong.
One likely problem was how I had asked it to practice. During its original training, the predictor started each step from an embedding of a real screenshot. If it made a slightly wrong prediction, the next example still began from the real next screenshot. Planning gives it no such reset: its second prediction starts from its first prediction, the third starts from its second, and so on. Small errors can carry forward until the model is working with embeddings unlike those of real screenshots.
CEM can make this worse. It searches through many sequences and favors whichever ones the model says get closest to a goal. If the model is especially wrong about one sequence, that mistake may make the sequence look unusually good to the search. We cannot identify the exact mistake from the failed plan alone, but a low predicted score clearly was not enough to trust it.
Rollout Fine-Tuning and the Second Attempt
To give the predictor practice with its own mistakes, I fine-tuned it on rollouts. A rollout starts from a real screenshot. The predictor guesses the next embedding, then uses that guess—not the embedding of the real next screenshot—to predict what follows the next action. This repeats across several recorded actions. We still have the real screenshots, so their embeddings can serve as targets for checking each guess. Training begins with short rollouts and gradually uses longer ones. The image encoder stays fixed so those targets do not move; only the predictor and action encoder are updated.
Step 12 MSE: original 0.4224 → fine-tuned 0.3045.
After fine-tuning, the first prediction actually got worse. But when the model had to keep predicting from its own guesses, the error grew more slowly. That was the tradeoff I cared about: a plan has to hold up for more than one button press.
I ran the planner again with the fine-tuned model. It found a sequence that selected Squirtle, even though it wasn’t the obvious route of pressing A twelve times. The emulator’s party count changed from zero to one, so this time the predicted success matched what happened in the game.
Starting state
Party count: 0
One successful plan still left me wondering how often this would work. I ran the planner with 100 fresh random seeds, keeping the model and starting state fixed, and tested every resulting plan in the emulator. 52 of the 100 plans acquired a starter, compared with zero for random button sequences and one for the same search using an untrained predictor. The learned model was helping, but it still failed almost half the time. And from this particular starting position, pressing A repeatedly already works. These results give me more confidence that the Squirtle run wasn’t just luck, while leaving plenty to test before I would trust the model to plan its way through more of the game.
The End…?
I never got to beat Blue (or Gary, depending on your preferences), but the model can do something. The final model trained end-to-end from scratch ended up being around ~12.5 million parameters. I think there would be more interesting versions of this which I may still try, like starting in Oak’s lab, having to walk to Oak, speak to him–therefore get through all his dialogue, and then choose a starter. This feels reasonable, but I am not sure how long of a rollout this would be, and I suspect that difficulty scales exponentially with plan length.
For now, I have Squirtle.
The implementation is in lePokeRed.