Technical Summary: Fast LeWorldModel

Problem Statement

Joint-Embedding Predictive Architectures (JEPAs), such as LeWorldModel (LeWM), have established a foundation for reconstruction-free visual world models capable of reward-free goal-conditioned planning. However, existing LeWM-style planners face two critical limitations due to their reliance on autoregressive one-step latent transitions:

1. Computational Inefficiency: To evaluate a candidate action sequence, the planner must repeatedly apply a local transition model step-by-step. This autoregressive rollout requires generating the entire imagined latent trajectory sequentially, repeatedly invoking action encoding and latent prediction modules. This process becomes a bottleneck during sampling-based planning (e.g., Cross-Entropy Method, CEM), where thousands of candidate sequences are evaluated.
2. Accumulated Prediction Error: Because each predicted state $\hat{z}_{t+k}$ depends on the previously imagined state $\hat{z}_{t+k-1}$, errors introduced at early or intermediate steps propagate and compound as the planning horizon increases. This makes long-horizon predictions increasingly unreliable.

Methodology: Fast LeWorldModel (Fast-LeWM)

The authors propose Fast-LeWorldModel (Fast-LeWM), a latent world model that replaces sequential autoregressive rollout with action-prefix prediction. Instead of predicting the next state from the current state, Fast-LeWM predicts the future latent state reached after executing a specific prefix of an action sequence, conditioned directly on the current observed latent state.

Core Architecture

The method consists of two primary components:

1. Action-Prefix Encoder:

   • Takes a candidate action sequence $(a_t, \dots, a_{t+H-1})$ and the current latent state $z_t$.
   • Maps $z_t$ to a "state token" and prepends it to the action tokens.
   • Uses a causal Transformer to process this sequence. Due to the causal mask, the output at position $k$ (the prefix token $p_{t,k}$) summarizes only the actions $(a_t, \dots, a_{t+k-1})$ and the current state context, preventing leakage from future actions.
   • This generates a set of prefix tokens $\{p_{t,1}, \dots, p_{t,H}\}$, where each token represents the accumulated effect of a specific action prefix.

2. Parallel Latent Predictor:

   • Takes the current anchor latent $z_t$ and the set of prefix tokens $\{p_{t,1}, \dots, p_{t,H}\}$.
   • Predicts all future latents $\{\hat{z}_{t+1}, \dots, \hat{z}_{t+H}\}$ in a single forward pass using a parallel predictor network.
   • The prediction for horizon $k$ is defined as $\hat{z}_{t+k} = G_\phi(z_t, p_{t,k})$.

Training Objective

Fast-LeWM employs dense prefix-level supervision. Unlike standard models that only supervise the final state or single-step transitions, Fast-LeWM minimizes the loss between predicted and ground-truth latents for every horizon $k$:
$$\mathcal{L}_{prefix} = \frac{1}{H} \sum_{k=1}^{H} \|\hat{z}_{t+k} - z_{t+k}\|_2^2$$
This forces the model to learn how states evolve continuously under different action prefixes, rather than just fitting local one-step transitions. A SIGReg regularizer is retained to prevent latent collapse.

Planning Protocol

During planning, Fast-LeWM maintains the same goal-conditioned CEM protocol as LeWM. However, instead of rolling out a sequence step-by-step, it encodes the action sequence once and queries all future horizons in parallel. An optional self-consistency term can be added to the cost function, penalizing discrepancies between a direct terminal prediction and a prediction derived via an intermediate prefix step, further stabilizing candidate selection.

Key Contributions

1. Identification of Bottlenecks: The paper identifies the local one-step transition interface in LeWM as the primary cause of slow autoregressive rollouts and error accumulation over long horizons.
2. Action-Prefix Reformulation: The authors propose reformulating latent dynamics modeling from single-step transitions to action-prefix prediction. This enables dense prefix-level supervision and transforms accumulated action effects into direct dynamic training targets.
3. Parallel Multi-Horizon Prediction: By decoupling future state predictions from sequential dependencies, the model achieves parallel state evolution prediction, significantly reducing computational cost and error propagation.

Experimental Results

The authors evaluated Fast-LeWM on four goal-conditioned planning tasks (Two-Room, Reacher, Push-T, OGBench-Cube) using the same protocols and datasets as LeWM.

• Planning Success: Fast-LeWM improved the average success rate from 85.8% (LeWM) to 90.5%. With the optional self-consistency term, this further increased to 92.0%.
• Efficiency:
  • Dynamics Module Time: Reduced from 31.4s to 8.0s (a 3.92× speedup) on the Two-Room task by replacing sequential rollout with parallel prediction.
  • Total CEM Solve Time: Reduced from 54.4s to 28.3s (a 48.0% reduction).
• Prediction Accuracy:
  • Open-Loop Error: Fast-LeWM achieved lower initial latent prediction error and a significantly slower growth rate of error as the horizon increased compared to LeWM.
  • Physical Probing: On the Push-T task, Fast-LeWM demonstrated superior performance in MLP probes for agent location, block location, and block angle, suggesting the learned latent space retains richer physical state information than LeWM or PLDM.

Significance and Claims

The paper claims that action-prefix prediction serves as an effective and efficient dynamics interface for latent world-model planning. By shifting from sequential autoregressive rollouts to parallel prefix-based queries, Fast-LeWM addresses the fundamental trade-off between planning speed and prediction reliability.

The authors emphasize that this approach does not require a new visual representation objective (staying within the JEPA framework) but rather optimizes how candidate action sequences are evaluated. The results suggest that modeling state evolution through accumulated action prefixes allows the dynamics model to learn more accurate, continuous state transitions, mitigating the compounding errors that typically plague long-horizon visual planning.