LeWAM improves visual control planning success and speed dramatically
LeWAM: End-to-End World and Action Modeling with JEPAs
RoboticsArtificial Intelligence
Summary
Planning how a robot should act by watching the world usually involves complex image processing. The team behind LeWAM developed a method that learns to predict future states and actions in a simplified, compact way without needing to reconstruct full images. This makes it much faster and more reliable at planning and controlling movements, especially in difficult or long tasks. Their approach works well both in simulations and on real robots, showing it can effectively guide actions.
What this means in practice
- •For robotics engineers: Design control algorithms that plan faster and succeed more often in robotic manipulation tasks using learned world and action models.
- •For simulation developers: Implement efficient behavior cloning to transfer policies from simulation to physical robots with improved accuracy on complex tasks.
Authors
Minghao Fu, Tavis Siebert, Eryk Halicki, Randall Balestriero
Abstract
Joint-embedding predictive architectures (JEPAs) learn compact latent world models without pixel reconstruction. Existing JEPA pipelines typically train visual representations and action-conditioned dynamics before adding a control policy or solver. We introduce LeWAM, an end-to-end JEPA for control with: (1) a joint training objective over next latent state prediction and action flow matching, and (2) an isotropic Gaussian constraint on the latent state, enabling a single model to support direct policy execution and efficient planning by evaluating sampled actions with its learned latent dynamics. Compared with prior JEPAs in visual planning, LeWAM improves success rate from 28.6% to 89.7% on contact-rich tasks and from 10.9% to 46.2% on long-horizon tasks, and plans 32.3 times faster. On behavior cloning tasks in simulation and physical robots, LeWAM is also competitive with established methods. Interpretability analyses show that LeWAM learns control-relevant representations and action distributions effectively guiding test-time planning and control.