Traj-LeWM: Path-Aware World-Model Planning via Latent Trajectory Cost

Xiaodi Huang; Ziyi Ding; Jingtian Wan; Yuchen Liu; Yuan Zhang; Xiao-Ping Zhang; Jiayu Chen; Zhang Zhang

cs.AIAug 14, 2026

Traj-LeWM: Path-Aware World-Model Planning via Latent Trajectory Cost

Authors: Xiaodi Huang, Ziyi Ding, Jingtian Wan, Yuchen Liu, Yuan Zhang, Xiao-Ping Zhang, Jiayu Chen, Zhang Zhang,

Organizations: Institute of Automation, Chinese Academy of Sciences · Shanghai Jiao Tong University · Tsinghua Shenzhen International Graduate School · INFIFORCE ·

Abstract

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 $\mathcal{E}_{\theta}$	Predictor $\mathcal{F}_{\phi}$	LTC $S_{\psi}$
$\mathcal{L}_{\mathrm{pred}}$	✓	✓
$\mathcal{R}_{\mathrm{sigreg}}$	✓
$\mathcal{L}_{\mathrm{path}}$	✓		✓
$\mathcal{L}_{\mathrm{mined}}$			✓

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 $\uparrow$	Target miss $\downarrow$	Collision $\downarrow$
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.

World Model-Based Planning

World model-based planning uses models of how environments change to compare imagined actions and choose goal-directed behavior.

See all 245 papers

Latent World Models Should Plan with Distances Along the Latent Manifold
Oct 18, 2026Zunzhi You, Linwei Tao, Siyu Xu +2Latent World Models Latent-Space Planning
The University of Sydney · Machine Learning Research, Apple
Fast LeWorldModel
Jun 24, 2026Yuntian Gao, Xiangyu XuLatent Dynamics Modeling World Model-Based Planning
Xi’an Jiaotong University
LeWAM: End-to-End World and Action Modeling with JEPAs
Oct 8, 2026Minghao Fu, Tavis Siebert, Eryk Halicki +1World Model Learning World Model-Based Planning
UC San Diego · ETH Zurich · University of British Columbia +1

Oct 18, 2026J/K move · Enter open · S save

Latent World Models Should Plan with Distances Along the Latent Manifold

Zunzhi You, Linwei Tao, Siyu Xu, Yunke Wang+1

The University of Sydney · Machine Learning Research, Apple

Latent World Models Latent-Space Planning

Open paper AskPDF arXiv

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.

Abstract

World Model-Based Planning

Explore similar work

Latent World Models Should Plan with Distances Along the Latent Manifold

Fast LeWorldModel

LeWAM: End-to-End World and Action Modeling with JEPAs