Planning through a memory of demonstrations
The standard Push-T test is easier than it looks
The caveat is that the standard test takes a demonstration, starts the robot partway through it, and sets the goal a fixed number of steps further along the same demonstration. That's a reasonable design, but in Push-T every demonstration ends with the block in the same place. So every test goal sits on a path that demonstrations already travel, heading for the spot they all end at. For 84% of the standard test's goals, the memory graph already held a nearly identical state from another demonstration.
The usual sample is also small. On the official 50 episodes, LeWM succeeded 68% of the time with the goal 50 steps away. On 250 fresh episodes drawn the same way, it succeeded only 44%. Changing nothing but the planner's random seed flipped 38 of 250 episodes between success and failure.
That's where I decided to build a harder test, where each task takes its start from one demonstration and its goal from another, so the robot has to travel between states that no single demonstration connects. It only counts as a success if the block stays on target for the last 10 steps. Under these rules the memory-graph planner succeeds 75% of the time on standard tasks and 45% on the new ones. The third and fourth methods I found work on that gap.
3. Fill gaps in memory with transformed memories
A memory graph is only as good as its coverage. When I deleted every stored state near the goal, success fell from 93% to 32%.
Push-T has a useful symmetry here. The physics of pushing a block doesn't change if you shift, rotate or mirror the whole table. So when the graph has no states near a goal, it can take demonstration segments from elsewhere, transform them to fit the gap, and plan through those. With these transformed segments, success came back to 85%. They also stretch a small memory a long way (a graph built from 1% of the demonstrations plus transformed copies did about as well as one built from 10%).
Filling a gap in memory with transformed demonstrations
a demonstration segment from elsewherethe same segment, shifted, rotated or mirrored
Block reached its target, 250 long-horizon tasks (goal 75 steps away)
93%full memory
32%states near the goal deleted
85%gap filled with rotated and mirrored copies
On the harder test, transformed segments raise success from 45% to 56%, and they help most where demonstrations are thinnest. When the goal lies in the direction the demonstrations flow, the memory-graph planner already succeeds 69% of the time. When it lies the opposite way, plain memory gets 27% and transformed memory gets 46%.
Harder test, by where the goal lies
about 83 tasks per row
69%
success rate
plain memorywith transformed segments
4. The cost was “wrong”
I also built an interactive demo where I could place the block anywhere and pick any goal. A goal has to be an image, and an image of the scene has to show the pusher somewhere, so I placed the pusher just behind the goal. In one run, the robot pushed the block to within 14 pixels of the target, close enough to count as solved. Then it pushed the block away from the goal, and kept moving the pusher toward where the pusher was in the goal image.
The goals in the harder test are frames from expert demonstrations, so each one also shows where the expert’s pusher happened to be at that moment. A planner that compares its prediction with the whole goal image is trying to put the pusher there too.
To measure how much this matters, I took the model out of the question. I ran the exact-physics planner twice on the same 50 tasks from the harder test, and the only difference was whether its cost included the pusher’s position as well as the block’s.
Block held on target, same 50 tasks
exact-physics planner
success with either cost30 succeed only without the pusher0 succeed only with the pusher included
Thirty tasks went from failure to success when the pusher was left out, and none went the other way. A split like that happens by chance about twice in a billion. The physics here is from the actual simulator, not predicted, so the 4x in success comes from how the goal was created.
The planner that scored the pusher did exactly what it was asked. On the 30 tasks that flipped, it got the pusher to within half a pixel of where the goal image put it, and lost the block on the way. Placing both the T and the pusher in the exact goal position is a much harder problem to solve.
Pusher included vs. not
Cost compares the block and the pusher
Cost compares the block only
target for the blockwhere the goal image put the pusher
But LeWM can’t simply be told to ignore the pusher. It scores a plan by the distance between two embeddings, and its embedding of a goal image holds the block and the pusher together, with no part labeled “pusher”. So, as a test, I measured how far a planner gets by searching over the pusher instead. It compares each prediction against 16 versions of the goal image, all with the block in the same place and the pusher moved to positions taken from stored demonstrations, and keeps the closest match. That doubled LeWM’s success on the harder test, from 11% to 23%, without retraining.
Sixteen versions of one goal
16 copies of the goal with the pusher moved, and the closest one sets the cost
This applies to any robot because a goal should just be a set of end states, not a copy of an image. “Put the cup on the shelf” is solved wherever the robot’s arm ends up, but a goal image picks one arm position and makes the planner match it. A good goal accepts every end state that solves the task, and still includes what the task needs from the robot, like coming to rest at the end.
The good news is that large sets of demonstrations already contain many real frames of the same object in the same place, each with the robot somewhere different. Scoring a plan against the closest of those goals allows for a larger set of valid end states to plan for.
What comes next
None of these four discoveries change the model's weights. They apply to LeWM as published, and they should carry over to larger JEPA world models that plan the same way, at no extra training cost.
The biggest open problem remaining is generalizing to tasks that don't have perfect demonstrations. On the harder Push-T test, the best planner holds the block on target 56% of the time, and most of the remaining failures are goals that run against the direction of the data. Closing that gap needs a higher level that can propose its own subgoals beyond the demonstrations, and do it better than today's learned hierarchies like H-JEPA [4], which so far match or slightly beat flat LeWM on Push-T.
Hierarchical subgoal planner on top of LeWM
Going in, I assumed LeWM's predictions would be the bottleneck, but a JEPA small enough to train in a few hours already predicts about as well as the simulator itself, and there are several gains that can be achieved simply by changing the way it plans. What it can't do yet is choose its own subgoals where no demonstration has been. That's what I'm currently working on, and if you want to see more, I'll post about it first on X, @haddyalchaer.
References
- [1] Y. LeCun, A Path Towards Autonomous Machine Intelligence. OpenReview, 2022.
- [2] L. Maes, Q. Le Lidec, D. Scieur, Y. LeCun, and R. Balestriero, LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels. arXiv:2603.19312, 2026.
- [3] C. Chi et al., Diffusion Policy: Visuomotor Policy Learning via Action Diffusion. RSS, 2023.
- [4] W. Zhang, B. Terver, M. Rabbat, Y. LeCun, and R. Balestriero, H-JEPA: End-to-End Learning of Hierarchical World Models for Visual Planning. arXiv:2610.06805, 2026.