TrekWM: Long-Horizon Planning via Shaped Latent Geometry and Adaptive Subgoals

Latent world models (LWMs) offer a promising framework for zero-shot visual control by planning in a latent representation space rather than raw pixels. However, latents trained exclusively for prediction do not naturally induce a metric suitable for evaluating planned trajectories: Euclidean distance between latents fails to reliably reflect true environment steps, leading to severe degradation over long task horizons. To address this, we propose TrekWM to improve long-horizon planning in LWMs via two complementary mechanisms. First, a ranking loss shapes the latent geometry so that latent Euclidean distance preserves temporal ordering when scoring candidate futures during planning. Second, we shorten the effective planning horizon using an adaptive subgoal proposer trained with a flexible temporal schedule, enabling the discovery of meaningful task landmarks rather than relying on rigid, fixed intervals. Together, these complementary improvements ensure both reliable local distance estimation and scalable global planning. Across eleven settings in five environments, TrekWM achieves the top performance in ten, outperforming competing methods by a substantial margin. The advantage is most pronounced on the longest-horizon tasks: $67.3\%$ vs. $0.7\%$ on Cube, $60.7\%$ vs. $23.3\%$ on PushT, and $100.0\%$ vs. $54.7\%$ on TwoRoom, in each case compared against the strongest baseline for that setting.