Source-linked AI summary
Fast LeWorldModel
Yuntian Gao, Xiangyu Xu
TL;DR
LeWM-style visual planning is costly because candidate sequences require autoregressive one-step latent rollouts that can accumulate prediction error over the horizon. Fast-LeWM replaces these rollouts with parallel action-prefix prediction, raising average success from 85.8% to 90.5% and reducing full CEM solve time by 48.0% (54.4s to 28.3s).
Problem
LeWM-style planning requires costly autoregressive one-step rollouts, with intermediate latent predictions reused and errors accumulating over longer horizons.
Method
Fast-LeWM predicts future latents directly from action-prefix representations in parallel, using dense prefix-level supervision instead of single-step transition modeling.
Results
Across all LeWM planning tasks, average success rose from 85.8% to 90.5%, while full CEM solve time fell 48.0% (54.4s to 28.3s).
Takeaways & Limitations
Action-prefix prediction provides an effective and efficient dynamics interface for latent world-model planning within the evaluated tasks.
Takeaways & Limitations
The study assumes an offline, reward-free dataset of observation-action trajectories.
Abstract
from arXiv · showhide
Joint-Embedding Predictive Architectures (JEPAs), including recent LeWorldModel (LeWM), have become a promising foundation for reconstruction-free visual world models. For visual planning, however, LeWM evaluates candidate action sequences by repeatedly applying a local one-step latent transition model. This autoregressive rollout makes planning computationally expensive and exposes the predicted trajectory to accumulated latent errors as the horizon grows. We propose Fast LeWorldModel (Fast-LeWM), a fast latent world model that replaces repeated local rollout with action-prefix prediction. Given the current latent and a candidate action sequence, Fast-LeWM encodes its prefixes and predicts the future latents reached after executing those prefixes in parallel. By making action prefixes the basic prediction unit, Fast-LeWM directly models action effects accumulated to different extents over multiple horizons. This prefix-level supervision forces the model to learn how states continuously evolve under different action prefixes, rather than only fitting one-step state transitions. During planning, the predictor can use the last prefix token from the encoded action sequence to evaluate the corresponding future latent without explicitly rolling through each intermediate imagined state. Across multiple tasks, Fast-LeWM improves average success over LeWM while substantially reducing planning time, achieving lower open-loop latent loss whose growth becomes significantly slower as the rollout horizon increases.
1 Introduction
Fast-LeWM addresses LeWM’s costly, error-prone autoregressive planning by directly predicting future latents for action prefixes in parallel. This reformulation enables dense prefix-level supervision and improves planning efficiency, success, and long-horizon prediction reliability.
- Motivation: LeWM evaluates candidate sequences through repeated one-step latent transitions, making planning slow and vulnerable to errors propagated from intermediate imagined states.The terminal prediction is obtained by composing local transitions, with each prediction depending on intermediate imagined latents.
- Fast-LeWM: Fast-LeWM replaces local rollout with action-prefix prediction, directly modeling future latents reached after executing prefixes of an encoded action sequence.Different prefixes represent different degrees of accumulated action effects and provide direct multi-horizon state-evolution targets.
- Method: A causal action-prefix encoder and parallel latent predictor map all prefixes to corresponding future latents in one forward pass.Each prefix token summarizes its preceding actions with the current latent context, while dense prefix-level supervision trains accumulated action effects.
- Planning benefits: Parallel prefix predictions avoid sequential dependence inside the predictor, reducing recursive rollout-error accumulation and accelerating state-evolution prediction.The planner predicts all corresponding future states in parallel rather than generating the imagined trajectory step by step.
- Results: 90.5% average success rate versus 85.8% for LeWM, with a 3.9× dynamics-module speedup and 48.0% lower full CEM solve time.The reported times are 31.4s to 8.0s for the dynamics module and 54.4s to 28.3s for the full CEM solve.
2 Related Work
Prior latent world models learn action-conditioned dynamics for planning in compact visual representations, while JEPAs offer reconstruction-free prediction of future embeddings. Fast-LeWM changes candidate-sequence evaluation by replacing LeWM’s sequential local rollout with action-prefix queries at multiple horizons.
- Latent World Models for Planning: Latent world models predict action consequences before execution, commonly encoding observations into compact latent states for planning or policy learning.PlaNet and Dreamer are cited as examples that learn latent dynamics from images and use imagined trajectories.
- Reconstruction-Free Visual Dynamics: JEPAs provide a reconstruction-free alternative to pixel reconstruction by predicting future or masked embeddings instead of pixels.The passage motivates this approach because pixel reconstruction may preserve visual details that are not necessarily relevant for control.
- Reconstruction-Free Visual Dynamics: LeWorldModel trains a compact end-to-end JEPA world model from raw pixels using next-embedding prediction and SIGReg-style Gaussian regularization.The next-embedding loss and Gaussian regularization are attributed to the cited JEPA and LeWM work.
- Fast-LeWM’s Evaluation Strategy: LeWM-style planning evaluates multi-step candidates by applying a local next-latent transition model sequentially, making later predictions depend on earlier predicted latents.This describes the prior evaluation strategy that Fast-LeWM modifies.
- Fast-LeWM’s Evaluation Strategy: Fast-LeWM instead uses action prefixes as multi-horizon queries, conditioning each future latent on the observed anchor latent and its corresponding action prefix.The supplied passage contrasts prefix-based evaluation with sequential local rollout, though the final word of the passage is truncated.
3 Method
Fast-LeWM is a reconstruction-free latent world model for reward-free, goal-conditioned planning that predicts future latents from the current latent and action prefixes rather than autoregressive one-step rollouts. This prefix-level interface enables parallel prediction, dense supervision across horizons, and faster candidate evaluation while retaining LeWM’s CEM planning objective.
- Problem setup: Fast-LeWM targets offline, reward-free goal-conditioned planning from pixel observations by learning a latent world model and comparing predicted latents with encoded goal observations.The current and goal observations are mapped into latent embeddings, and planning searches over action sequences in latent space.
- Autoregressive baseline: LeWM-style dynamics evaluate candidate sequences through autoregressive one-step rollouts, requiring repeated generation of intermediate imagined latents and increasing planning cost with many CEM candidates.The sequential chain also introduces intermediate predicted latents early into the rollout, exposing later predictions to accumulated latent errors.
- Prefix-level prediction: Fast-LeWM predicts each future latent directly from the observed anchor latent and its corresponding action-prefix token, allowing different horizons to be queried independently.For horizon k, the predictor uses ˆzt+k = Gϕ(zt, pt,k), where pt,k summarizes the accumulated effect of the first k actions.
- Training and computation: Dense prefix-level supervision trains intermediate and terminal latent predictions, while all queried horizons share one encoding pass and one parallel prediction pass during planning.This objective teaches how latent states evolve as actions accumulate rather than only fitting local one-step transitions.
- Action-prefix encoder: A causal Transformer action-prefix encoder produces horizon-specific tokens that attend only to the actions within each prefix, while a prepended state token supplies current-scene context.The state token helps distinguish identical open-loop actions whose effects differ across scene configurations.
- Planning objective and results: 90.5% average success rate is achieved by Fast-LeWM versus 85.8% for LeWM, increasing to 92.0% with the optional self-consistency term.Fast-LeWM retains LeWM’s CEM objective, while the consistency penalty favors candidates with stable terminal predictions under different prefix decompositions.
4 Experiments
Across matched planning settings, Fast-LeWM improves success while reducing computational cost through parallel action-prefix prediction. Its prefix-based dynamics also slows open-loop latent error growth, preserves physical information, and outperforms simpler speed-oriented ablations.
- Planning efficiency: Parallel action-prefix encoding and latent prediction substantially reduce planning cost relative to LeWM’s sequential predictor under the same CEM budget.Fast-LeWM reports both dynamics-evaluation time and full CEM solve time; the latter also includes goal- and observation-related costs.
- Open-loop prediction: Fast-LeWM consistently lowers open-loop latent prediction error and slows its growth as the trajectory evolves, mitigating error accumulation over longer horizons.At t = 25, Fast-LeWM makes one maximum-horizon prediction, whereas LeWM requires five autoregressive local transitions; at t = 50, the corresponding counts are two and ten.
- Open-loop prediction: Decoded Fast-LeWM rollouts remain more consistent with the target trajectory, while LeWM shows visible drift at longer horizons from accumulated rollout error.Both models use the same initial observation and future action sequence in the qualitative comparison.
- Latent probing: Under MLP probes on PushT, Fast-LeWM achieves the lowest MSE and highest correlation for agent location, block location, and block angle.Its linear-probe performance is comparable to LeWM and consistently better than PLDM, indicating preserved linearly accessible physical structure.
- Ablations: Long-Action LeWM performs poorly, showing that extending each next-state transition to a longer action block does not by itself provide fast planning.Removing dense prefix supervision creates a terminal-only variant, while conditioning the prefix encoder on the current state token improves performance with only a small overhead.
5 Conclusion
Fast-LeWM is a reward-free visual world model that predicts future latents directly from action-prefix representations rather than autoregressive one-step rollouts. This enables parallel multi-horizon prediction, reduces repeated dynamics evaluations and latent error accumulation, and improves planning efficiency and success across four goal-conditioned tasks.
- 5 Conclusion: Fast-LeWM predicts future latents directly from the current latent and action-prefix representations, enabling parallel multi-horizon prediction.It replaces autoregressive one-step rollout with action-prefix prediction.
- 5 Conclusion: The action-prefix design reduces repeated dynamics evaluations, mitigates latent error accumulation, and provides dense supervision over continuous state evolution.These benefits follow from making action-prefix representations the basis for prediction rather than repeatedly applying one-step transitions.
- 5 Conclusion: Across four goal-conditioned planning tasks, Fast-LeWM improves planning success while substantially reducing CEM solve time and open-loop prediction cost.The experiments evaluate the model in reward-free visual planning.