LeMario: Training a JEPA World Model on Super Mario Bros
benjamin-bai.com133 points by kevinjosethomas 2 days ago
133 points by kevinjosethomas 2 days ago
The author hints at this but it seems like one issues is that while JEPA is good at distinguishing between unpredictable noise and predictable features, the model has no way of assigning importance to different predictable features.
So for a system where it’s very difficult to exactly reach the desired end state, the model needs to choose between (for example):
- reaching a relatively achievable scene where 95% of the features in the latent are correct, which includes stuff like visible enemies, Mario’s position on the screen etc
- reaching a far more difficult to access scene where there’s a bunch of differences in the actual level visuals, but theres a match on the latent for the tiny set of pixels in HUD that indicate you’ve hit the victory condition
We obviously know that it’s not good enough to reach an early scene that looks similar to the victory condition but isn’t. The model doesn’t.
In a sense, this is what the linear probe helps with - it allows us to re-weight the latent and say “actually, while the latent encodes many things about the world, the thing we really care about is the X position.”
I’d be curious what happened if rather than planning actions on cross entropy of a final scene, the model just tried to find the actions that maximize the predicted X value of the probe.
Hey! Author here! Yes, I completely agree with your take! I want to run another experiment around whether the JEPA can focus its distance metric on the task-relevant features. But I think this is also where world modeling starts to blend into RL or goal-conditioned learning, the world model can learn what is predictable, but an external objective still has to tell it what matters. The position probe was essentially a lightweight way of adding that task-specific weighting without retraining the representation.
Also I actually did try the experiment where we maximized the predicted X! It worked across open ground, moving Mario from x=40 to roughly x=360. But... it broke down at the first obstacle, it repeatedly chose to jump in place. The model predicted those jumps would increase X, but Mario’s real position stayed around x=342. It just couldn't get over the error in its learned jumping dynamics.
Thanks a bunch for reading :)))
I think JEPA is super interesting, but I feel like this example highlights some of the challenges of long horizon planning. For one, chunking the planning stage into a bunch of intermediate goals seems really limiting, because a lot of what makes model based control interesting is that we don’t want to impose a solution strategy (because we want to solve problems we don’t know how to solve)
Another thing that has been bothering me is that you have to write the goal in input space. That doesn’t align with all problems, for some problems there could be many different states that satisfy a goal. For Mario maybe it’s ok, but there’s some weirdness still, like should the goal state be Mario at the finish line of the level with a specific timer state in the frame header? What about optimizing the number of points?
Also it’s interesting to think about how you would get Mario to reliably jump on koopas and goombas. IIUC JEPA models are usually trained with random rollouts, and then you’d handle this sort of intermediate goal in the planning optimizer? But that seems inefficient, and including some planning in the pretraining rollouts might be necessary to get enough relevant intermediate states. And then it starts feeling like reinforcement learning…
I’d be happy to have a check on my intuition here, or pointers to interesting writing on these topics
p.s. on topic, I liked the debugging strategies used in the blog post, that was my favorite part of the writeup
Author here! Thanks for reading!
Your intuition is very close to what I took away from the project as well! The intermediate image goals were intended as a diagnostic rather than a general solution. The checkpoints helped Mario move farther, but yeah like you said they also imposed part of the solution and didn’t fix the underlying representation problem. The goal is definitely to have a more general system, which would need to discover useful subgoals itself.
I also agree about goals in input space. JEPA itself doesn’t require this, but the LeWorldModel planner uses the embedding of one goal image as its objective. Which discussed, Mario can reach the correct location but have a different timer, animation, or enemy state and still be considered far away. Ideally the objective would represent a set of successful states or only task-relevant features. But that would require explicitly introducing something like a reward, goal classifier, or learned value function (like the probe), which would no longer be reward-free.
When it comes to killing the enemies, the world model can only predict transitions that its dataset covers well. And to clarify the JEPA objective doesn’t specifically require random rollouts, it can learn from random, expert (like they did in the original paper for Push-T). The planner can search over known dynamics, but it can’t reliably invent dynamics the model never learned. Exploration, demonstrations, or online data collection would help, and I agree that this starts to blur into model-based RL.
My current view is that JEPA can provide the representation and predictive dynamics, but it doesn’t automatically solve exploration, goal specification, or long-horizon control. Dreamer and TD-MPC are relevant examples of combining learned world models with value learning, while goal-conditioned and hierarchical RL address goal sets and learned subgoals.
Thanks so much for the thoughtful comment :D
Such a gem, thanks to the author for sharing it's findings :)
The only problem I have with planing in latent space is that it can be really noisy and not representative of the positions in the game (the latent are trained for semantic, so the optimizer can focus a set of specific features and can skip positions, which means it cannot know "where" to go by optimizing on the latents directly).
Hey, author here! Thanks so much for reading :) I totally agree, the latent captures details useful for prediction, but not necessarily for control, making planning noisy. Still, I was surprised by how well it understood horizontal position after less than two hours of training on one A100!
How big is this model? Is it expensive to train? Can it successfully complete a small level? Thanks.
Nice job, i have also been working on a few JEPA based models during the last few months. Trying to make more efficient LLMs.
I feel like you hit the main issues in the use of jepa models (well except collapse but sigREG more or less solves the collapse issue).
The main issues in JEPA as i see it is pushing the latent space toward representing features that are needed for good planing. A thing which is especially a problem in hierarchical planing.
You prime a JEPA world model to predict changes based on actions but you never really push it to use those actions. You simply hope that it will use them. If your latent is big enough and the actions effect on the world is simple enough it tends to work out but those qualifiers are not always small things.
Secondarily finding actions for the higher level JEPA Predictors.
LeWorldModel encodes multiple movements in to higher level actions. But this is a not a very good idea. It solves a basic issue with planing where the predictions degrade after a set nr of steps. But it does not solve the issue of higher level actions not actually being button presses.
The higher level actions for your mario game version would be things like: get the coin, Kill an enemy or get to the end of this stage.
You cant just encode many button presses in to those types of things. You need to discover those actions somehow.
Am training a JEPA inspired model on my entire system.
BPF programs collate process lists, memory, cpu, network, storage and video buffer/frames into a model
I can then browse it with a custom Vulkan powered browser and recreate the observed states in a 3D space.
The fundamentals work. I recently moved into runtime testing sussing out edge cases.
Aside from the nonsensical code to interface with the mess that is Linux, the model and browser code is clean geometric transformation.
Rather than save the context from the filesystem I label snapshots of text files in the browser. All text is treated like geometric data.
Here's my two cents as a mere paper reader;
JEPA is really just a generalized encoder, so the JEPA created latents should be fed into a transformer trained with either user data or RL policy.
Although the above might not work great either because you said that vertical position was not predicted well!
Great article and I hope that you can carry on with the JEPA research
Really enjoyed this.
But I believe the goal isn’t a place - some absolute location to the right. The goal is an action: to always be in the state of holding the right d-pad button down, or taking some intermediate action so we can go back to holding the right d-pad down again.
benbye - do you have any intuition why the model only needs 4 frames to work? I played with LeWorldModel and I am mystified why the latent prediction frame number is so small. It's an almost Markov like property where a short history is packaged as state.
man I'm so brainrotted, I just see these names and I laugh
[dead]