LeMario: Training a JEPA World Model on Super Mario Bros

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.

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

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.

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).

man I'm so brainrotted, I just see these names and I laugh