Reproduce “LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC” on OpenResearch

About this paper

Shashank Hegde, Alexander Popov, Elie Aljalbout, Nikolai Smolyanskiy

LeWAM combines action generation and future-state prediction in a decoder-free world model for pixel-based simulated manipulation. Its planner searches sequences of policy noise inputs, rather than arbitrary actions, and scores imagined outcomes against a demonstrated goal. Across six distractor-augmented tasks, noise-space planning improves closed-loop success, while action-space MPC yields only 3–8% for every model. Evaluations are simulated, not real-robot demonstrations.

Key insights

• Across six tasks, LeWAM’s latent achieved average task-state probe R² of 0.80 and distractor R² near 0.00; the reconstruction-based RecWAM had distractor R² of 0.73.
• Constraining planning to the policy’s generative process can reduce the planner’s ability to exploit dynamics-model errors; directly optimized action-space MPC achieved only about 3–8% success in the reported evaluations.
• Keeping noise proposals near the prior’s typical norm helped on some tasks: with projection, DS-CEM reached 78.8% on Can versus 23.0% without projection.

Summary by alphaXiv.