Organizations: Institute of Automation, Chinese Academy of Sciences · Shanghai Jiao Tong University · Tsinghua Shenzhen International Graduate School · INFIFORCE · University of Science and Technology of China · The University of Hong Kong · Peking University
LeWM is a lightweight visual world model that learns latent dynamics end-to-end from pixels and ranks candidate action sequences by the distance between their predicted endpoints and the goal. However, LeWM has two limitations. First, during training, it learns local next-step transitions without evaluating complete trajectories relative to the task goal. Second, during planning, it ranks candidates solely by predicted endpoint distance. Because model predictions may differ from actual execution outcomes, the candidate whose predicted endpoint is closest to the goal may not perform best when executed in the environment. The evolution of the complete predicted trajectory can therefore provide complementary information beyond endpoint distance. To address these limitations, we propose Traj-LeWM, which retains LeWM's local-dynamics objective and endpoint score while introducing a goal-conditioned latent trajectory cost (LTC) that aggregates trajectory-level information as a complementary signal. During training, LTC-based trajectory-preference supervision complements next-step prediction in shaping the shared representation. During planning, LTC is combined with endpoint distance to incorporate intermediate-path information into candidate ranking. With joint endpoint-plus-LTC scoring, Traj-LeWM outperforms LeWM on Push-T, OGBench-Cube, Reacher, and Two-Room by 3, 14, 7, and 7 percentage points, respectively. Controlled experiments and ablations further verify the complementary roles of trajectory-level representation shaping and path-aware candidate ranking.
Figures & tables
Figure 1: Overview of Traj-LeWM. Stage 1 encodes the observation trajectory and goal into latent representations used by LTC. Stage 2 trains LTC with synthetic and mined closed-loop trajectory preferences while retaining LeWM’s next-step prediction objective. Stage 3 ranks CEM candidates using the joint endpoint-plus-LTC score.
Loss
Encoder Eθ
Predictor Fϕ
LTC Sψ
Lpred
✓
✓
Rsigreg
✓
Lpath
✓
✓
Lmined
✓
Table 1: Direct gradient flow from each loss term.
Model
Push-T
Cube
Reacher
Two-Room
GCBC
75
84
–
100
GCIVL
33
56
–
100
GCIQL
20
64
–
100
PLDM
78
65
78
97
DINO-WM
74
86
79
100
LeWM
96
74
86
87
Table 2: Closed-loop planning success rates (%) across four environments. For Traj-LeWM, results are averaged over three evaluation seeds and rounded to the nearest percentage point. Bold indicates the best result for each task.
Figure 2: Controlled ranking tests. (a) Endpoint-only AUC change from LeWM to Traj-LeWM. (b) Correct-ordering change after adding LTC for candidates with large prediction–execution discrepancies.
Figure 3: Endpoint-controlled tests on Reacher. (a) LTC correctly ranks 94.4% of 207 endpoint-matched success–failure pairs. (b) Intermediate-state perturbations increase LTC cost in all 200 pairs with fixed endpoints.
Figure 4: Real-robot setup and two example start–goal tasks. Each row shows the initial observation, an intermediate state, and the visual goal 1.5 seconds later.
Model
Success ↑
Target miss ↓
Collision ↓
LeWM
10/20 (50%)
6/20 (30%)
4/20 (20%)
Traj-LeWM
14/20 (70%)
4/20 (20%)
2/20 (10%)
Table 3: Real-robot outcomes for LeWM and Traj-LeWM on the same 20 start–goal tasks using a Franka FR3.
Training
Score
Push-T
Cube
Reacher
Two-Room
Avg.
LeWM
Endpoint
96
74
86
87
85.8
Traj-LeWM-syn
Endpoint
96
79
80
94
87.3
Traj-LeWM
Endpoint
98
80
84
94
89.0
Traj-LeWM
LTC
78
69
94
80
80.3
Traj-LeWM
Joint
99
88
93
94
93.5
Table 4: Ablation of training supervision and planning score (success rate, %). Bold indicates the best result for each task.
Figure 5: Final-layer CLS-to-patch attention along a Two-Room trajectory. Red indicates higher attention in Traj-LeWM/LeWM.
Figure 6: Agreement between candidate selectors on shared candidate pools. Panel (a) reports exact top-1 agreement, and panel (b) reports the fraction of shared candidates in the two top-5 sets. Bars are query means, error bars are query-level bootstrap 95% confidence intervals, and the gray band denotes the expectation for two independent rankings.
Figure 7: Paired change in variance-normalized open-loop latent rollout error. Negative values favor Traj-LeWM; shading denotes paired-bootstrap 95% confidence intervals. The dashed line marks the five-step planning horizon.
Latent world models such as LeWM plan by searching for an action sequence whose predicted terminal latent is closest to the goal latent under an ambient norm, such as the Euclidean distance. Planning with such a cost assumes that proximity in the ambient latent space reflects proximity to the goal in the environment. Studying LeWM, we find that this assumption holds only in a small neighborhood around the goal, whereas farther away the L2 cost forms spurious low-cost basins that trap the planner. Yet neighborhood-based embeddings show that the latents retain the underlying physical state. We therefore hypothesize that the latents lie near a curved, low-dimensional manifold, and that the terminal cost should measure the distance along this latent manifold rather than through the ambient space. To estimate this manifold distance, we propose the Neighborhood Graph Cost (NGC), which computes shortest-path distances on a mutual nearest-neighbor graph constructed over training latents. NGC changes only the terminal cost and requires no retraining. This single change improves the planning success rate of LeWM across all four benchmark tasks and nearly solves TwoRoom and Reacher: even at the farthest goal offsets tested, NGC reaches 95% and 100% success, whereas the L2 cost falls to 27% and 81%. These results support the manifold hypothesis and our argument that latent world models should plan with distances along the latent manifold.
Zunzhi You, Linwei Tao, Siyu Xu +2
The University of Sydney · Machine Learning Research, Apple
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 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.
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.
Minghao Fu, Tavis Siebert, Eryk Halicki +1
UC San Diego · ETH Zurich · University of British Columbia +1