|
|
|
|
|
|
by rsfern
16 days ago
|
|
|
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.
Thanks so much for the thoughtful comment :D