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