Bidirectional transformer improves robot action prediction and planning
LeWAM: A JEPA World Action Model with Diffusion-Steering-Based MPC
RoboticsArtificial IntelligenceComputer Vision and Pattern Recognition
Summary
Predicting robot actions and what will happen next is tricky because traditional methods often rely on messy visual details that confuse the predictions. The authors introduce LeWAM, a special kind of AI model that looks both forward and backward in time to better understand actions and policies without relying on visual reconstructions. This new model helps predict robot and object states more accurately, ignores distracting visual noise, performs actions similarly to existing methods, and improves planning by focusing on the action planning space rather than raw actions. Overall, LeWAM supports more reliable robot control and planning.
What this means in practice
- •For robotics engineers: Improve robot state prediction and control by using LeWAM's bidirectional model for more accurate action and policy forecasting.
- •For control systems developers: Enhance model predictive control by planning in the policy noise space to achieve better closed-loop performance under model inaccuracies.
Authors
Shashank Hegde, Alexander Popov, Elie Aljalbout, Nikolai Smolyanskiy
Abstract
World action models (WAMs) predict actions and future observations, typically from a reconstruction-based representation that carries noisy, redundant information which can complicate downstream predictions. We introduce LeWAM, a bidirectional transformer for forward, backward, inverse dynamics and policy prediction, on a decoder-free JEPA latent trained end-to-end through all four modes. We see the following benefits: 1) Alignment: linear probes read robot and object state from LeWAM's latent better than from a regular Le World Model (a forward-only JEPA world model), while the latent ignores visual distractors as well as LeWM does and far better than a reconstruction-based WAM. 2) Acting: Closed-loop evaluations of LeWAM match a regular flow-matching policy trained on the same encoder at matched size, while also providing a world model. 3) Planning: Sampling raw actions when planning with WAMs lets MPC exploit dynamics-model inaccuracies; planning in the noise space of the policy head instead improves the closed-loop performance of these WAMs.