Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning
About
Recently, world models have emerged as a promising paradigm for building intelligent agents by learning predictive models that estimate future environment states conditioned on observations and actions. In particular, JEPA-style latent world models provide an efficient alternative to pixel space prediction by learning action-conditioned dynamics in compact representation spaces. However, existing latent world models typically rely on one-step prediction and must be recursively rolled out for long-horizon planning, which leads to compounding errors and a mismatch between training objectives and downstream planning tasks. To address this limitation, we propose Variable-length Latent World Models (VLWMs), a framework that learns to predict future latent states conditioned on action sequences of variable lengths. Instead of training only on one-step transitions, VLWMs directly model temporally extended dynamics, allowing the same predictor to evaluate action plans over different horizons. We further introduce a curriculum training strategy that progressively expands the action horizon, stabilizing optimization from short-range dynamics to long-range prediction. At test time, we design planning methods tailored to VLWMs to better exploit their variable-length predictive capabilities. Experiments on long-horizon control tasks show that VLWMs significantly improve latent space world models, achieving 13\% average improvement over the state-of-the-art LeWM across different datasets, with especially large gains on tasks requiring extended planning. These results suggest that VLWM provides a simple yet effective paradigm for improving long-horizon prediction and planning in latent world models.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Goal-reaching Planning | PushT Long-Horizon (50 steps) | Success Rate60 | 16 | |
| Goal-conditioned Planning | PushT offset 25 | Success Rate94 | 6 | |
| Goal-conditioned Planning | PushT offset 75 | Success Rate20 | 6 | |
| Goal-conditioned Planning | PushT offset 100 | Success Rate12 | 6 | |
| Goal-conditioned Planning | OGBench-Cube offset 50 | Success Rate54 | 6 | |
| Goal-conditioned Planning | OGBench-Cube offset 75 | Success Rate54 | 6 | |
| Goal-conditioned Planning | OGBench-Cube offset 100 | Success Rate50 | 6 | |
| Goal-conditioned Planning | TwoRoom offset 25 | Success Rate96 | 6 | |
| Goal-conditioned Planning | TwoRoom offset 50 | Success Rate80 | 6 | |
| Goal-conditioned Planning | TwoRoom offset 75 | Success Rate70 | 6 |