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

2026-08-14

AI Summary

This study addresses the inaccurate trajectory evaluation in lightweight world models caused by exclusive reliance on terminal distance. We propose Traj-LeWM, which introduces a goal-conditioned latent trajectory cost as a complementary signal. During training, this method optimizes representations via trajectory preference supervision, while during planning, it integrates path information to refine candidate ranking, thereby achieving path-aware world model planning. Through the dual mechanisms of trajectory-level representation shaping and path-aware ranking, Traj-LeWM effectively enhances planning precision. Experiments across four benchmarks, including Push-T, demonstrate performance improvements of 3%, 14%, 7%, and 7%, respectively, validating the efficacy of the proposed approach.

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.

Problem

• Visual World Model
• Trajectory Planning
• Latent Dynamics
• Candidate Ranking
• Goal-Conditioned Learning

Innovation

• Latent Trajectory Cost
• Path-Aware Planning
• Visual World Model
• Trajectory-Preference Supervision
• Joint Scoring