🤖 machine learning

Beyond the Next Step: Variable-Length Latent World Models for Long-Horizon Planning

This paper introduces Variable-length Latent World Models (VLWMs), a framework that learns to predict future latent states over variable action horizons through curriculum training, thereby overcoming the compounding errors of traditional one-step models and significantly improving long-horizon planning performance.

Original authors: Tianqi Du, Qi Zhang, Yifei Wang, Yisen Wang

Published 2026-06-23

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tianqi Du, Qi Zhang, Yifei Wang, Yisen Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot how to navigate a maze or push a block to a specific spot. To do this, the robot needs a "mental map" or a World Model. This model acts like a simulator in the robot's brain, allowing it to imagine "What happens if I move my arm this way?" before actually doing it.

For a long time, these mental simulators had a major flaw: they were like a person taking a single step, looking at where they landed, taking another single step, looking again, and so on. If the robot made a tiny mistake in its first step, that error would get bigger and bigger with every subsequent step. By the time it tried to plan a long journey, its mental map would be completely wrong. This is called compounding error.

The paper you provided introduces a new system called VLWM (Variable-length Latent World Models). Here is how it works, explained simply:

1. The Old Way: The "Baby Steps" Problem

Think of the old models (like the one they call LeWM) as a student learning to walk by only practicing one step at a time.

The Training: The teacher only asks, "If you take one step, where will you be?"
The Problem: When the student tries to plan a 100-step journey, they have to imagine taking 100 single steps in a row. If they stumble on step 1, their prediction for step 100 is wildly off. They are great at short trips but terrible at long ones.

2. The New Way: The "Leapfrog" Method (VLWM)

The authors propose a smarter way to train the robot's brain. Instead of just practicing one step, they teach the robot to predict the future at any distance.

Variable-Length Training: Imagine a teacher asking the student: "If you take 3 steps, where will you be?" Then, "If you take 5 steps, where will you be?" Then, "If you take 10 steps?"
The Result: The robot learns to "leapfrog" over the middle ground. It doesn't have to calculate every tiny movement to get to the end; it can jump straight to the destination of a 5-step or 10-step sequence. This prevents the small errors from piling up.

3. The "Action-as-Token" Trick

In the old models, the robot's actions (like "move left" or "push forward") were hidden inside the math in a rigid way. It was like trying to read a book where the words were glued together.

The New Trick: The authors treat every action as a distinct token (like a word in a sentence). They mix the "state" (where the robot is) and the "actions" (what it does) into one long sequence, just like a sentence.
Why it helps: This allows the robot to read a sentence of 3 actions or a sentence of 10 actions without changing its brain structure. It's flexible and natural.

4. The "Climbing the Mountain" Strategy (Curriculum Learning)

You can't teach a baby to run a marathon on day one. If you try to teach the robot to predict 100 steps immediately, it gets confused and fails.

The Solution: They use a Curriculum.
- Stage 1: The robot only learns to predict 1 step.
- Stage 2: It learns to predict 1 or 2 steps.
- Stage 3: It learns 1, 2, 3, or 4 steps.
- Final Stage: It can predict any distance up to the maximum.
This builds a strong foundation of short-term skills before adding the complexity of long-term planning.

5. Planning: The "Chunking" Strategy

When the robot actually needs to solve a problem (like navigating a maze), it doesn't just use one method. It has a toolbox of strategies:

Fixed Jump: It decides to always jump 5 steps at a time.
Random Jumps: It jumps 2 steps, then 7, then 3, randomly.
Long-to-Short: It starts with big jumps to get a general idea of the path, then switches to tiny, precise jumps as it gets closer to the goal.
The Benefit: Because the robot was trained to handle any jump size, it can mix and match these strategies on the fly to avoid getting stuck.

The Results: What Did They Find?

The authors tested this on three different tasks:

PushT: Pushing a T-shaped block.
OGBench-Cube: A robot arm picking up a cube.
TwoRoom: A robot navigating from one room to another through a door.

The Findings:

Short Trips: The new system performed just as well as the old ones.
Long Trips: The new system was significantly better. In the "TwoRoom" task, when the goal was far away, the old system succeeded only 36% of the time, while the new system succeeded 68% of the time.
Why? The old system got lost because of the "compounding errors" of taking one step at a time. The new system stayed on track because it learned to see the bigger picture.

Summary

The paper argues that to make robots good at long-term planning, we shouldn't just teach them to take one step at a time. Instead, we should teach them to predict the future in chunks of varying sizes, starting small and getting bigger. This simple change allows them to plan much further ahead without losing their way.

Technical Summary: Variable-Length Latent World Models for Long-Horizon Planning

1. Problem Statement

World models have emerged as a paradigm for building intelligent agents by learning predictive models of environment dynamics. While recent latent world models (e.g., JEPA-style architectures) offer efficient alternatives to pixel-space prediction by operating in compact representation spaces, they face significant limitations in long-horizon planning.

The primary issues identified are:

Compounding Errors: Existing models typically rely on one-step prediction ( $t \to t+1$ ). For long-horizon tasks, these models must be recursively rolled out, causing prediction errors to accumulate rapidly, leading to temporal drift and degraded trajectory quality.
Training-Planning Mismatch: Most models are trained with a fixed-step prediction objective. This provides supervision only at a local temporal scale, failing to explicitly optimize for long-range consistency. Consequently, the training signal does not align with the nature of long-horizon planning problems.
Architectural Rigidity: Current latent world models often incorporate actions via dedicated conditioning mechanisms (e.g., FiLM layers or cross-attention) that fuse actions into hidden states. This design tightly couples the model to a fixed action layout (usually one action per step), making it difficult to disentangle the contributions of multiple actions or handle variable-length action sequences naturally.

2. Methodology: Variable-Length Latent World Models (VLWM)

The authors propose VLWM, a framework designed to learn latent environment dynamics across multiple prediction horizons within a single model.

2.1 Variable-Length Supervision

Instead of restricting supervision to the immediate next step, VLWM trains a single predictor $g_\phi$ to forecast the latent state at an arbitrary future offset $k \ge 1$ . The training objective is:
$\mathcal{L}_{VLWM} = \mathbb{E}_{t, k \sim p(k)} \left[ \| g_\phi(s_t, a_{t:t+k-1}) - \bar{g}(\bar{f}_{\bar{\theta}}(o_{t+k})) \|_2^2 \right]$
where $s_t$ is the current latent state, $a_{t:t+k-1}$ is a variable-length action segment of length $k$ , and the target is the latent state at time $t+k$ . This unifies short-term dynamics and long-range semantic consistency under a single objective.

2.2 Action-as-Token Formulation

To natively support variable-length action sequences, VLWM departs from layer-internal conditioning. Instead, it treats action embeddings as first-class transformer tokens.

Input Sequence: The model receives a sequence $X$ containing a history of state tokens followed by a variable number of action tokens ( $e_t, \dots, e_{t+k-1}$ ).
Mechanism: State and action tokens are distinguished by additive type embeddings. The transformer is causal, and the predicted latent $\hat{s}_{t+k}$ is read directly from the output of the last token in the sequence.
Benefit: This eliminates the need for predefined action-conditioning interfaces, allowing the model to condition on an arbitrary number of past actions without architectural changes.

2.3 Curriculum Training Strategy

Directly training with large prediction horizons from the start is unstable due to noisy gradients and optimization failure. VLWM employs a curriculum learning strategy:

Training is divided into stages. In stage $j$ , the horizon $k$ is sampled uniformly from $\{1, 2, \dots, j\}$ .
The model starts with $k=1$ and progressively expands the range of sampled horizons up to a maximum $K_{max}$ .
This ensures that long-range semantics are built upon well-learned short-horizon dynamics, stabilizing optimization.

2.4 Planning with Variable-Length Jumps

During inference, the planning horizon $H$ is partitioned into contiguous chunks $H = k_1 + k_2 + \dots + k_M$ . The planner evaluates the rollout as a sequence of variable-length latent jumps rather than recursive one-step predictions. The paper proposes three planning strategies:

P1 (Fixed but Task-Adaptive): Uses a constant chunk size $k^*$ selected per task (small for fine-grained control, large for long-horizon goals).
P2 (Random-Horizon Sampling): Samples chunk sizes uniformly from $\{1, \dots, K_{max}\}$ at each boundary.
P3 (Long-to-Short): Starts with large chunks to cover long-range structure and progressively shrinks chunk sizes toward the end to refine fine-grained dynamics (a coarse-to-fine search).

3. Key Contributions

Unified Framework: Introduction of VLWM, a framework that learns latent dynamics across multiple horizons and naturally supports variable-length action sequences via an "action-as-token" formulation.
Stable Training & Planning: Proposal of a scalable curriculum learning strategy for stable variable-length training and a suite of planning algorithms that adaptively leverage different prediction horizons within a single model.
Empirical Validation: Demonstration that VLWM consistently outperforms fixed-horizon world models (specifically LeWM) on long-horizon planning benchmarks, with analyses showing improved long-range predictability and downstream controllability.

4. Experimental Results

The authors evaluated VLWM on three goal-conditioned planning environments: PushT (2D manipulation), OGBench-Cube (3D robotic manipulation), and TwoRoom (2D navigation).

Performance Gains: VLWM achieved an average 13% improvement over the state-of-the-art LeWM across datasets. The gains were most significant in long-horizon regimes (e.g., on TwoRoom with a horizon of 100 steps, success rate improved from 36% with LeWM to 68% with VLWM).
Error Mitigation: Visualizations showed that while LeWM's recursive rollouts drifted and failed to cross environmental boundaries (e.g., doors), VLWM's chunked variable-length jumps stayed close to ground-truth trajectories.
Representation Quality: Linear probe evaluations on frozen encoders revealed that VLWM produced latents with higher Pearson correlation for physical properties (e.g., block angle) compared to LeWM, suggesting better capture of slow semantic state changes.
Curriculum Necessity: Ablation studies confirmed that the curriculum training strategy is essential for stable variable-length training, particularly on tasks requiring extended planning.

5. Significance and Claims

The paper claims that VLWM provides a simple yet effective paradigm for improving long-horizon prediction and planning in latent world models.

Beyond Engineering: The authors argue that the struggle with long-horizon planning is not merely an engineering issue but a consequence of training objectives. By shifting from fixed-step to variable-length supervision, VLWM aligns the training signal with the requirements of long-horizon reasoning.
Unified Interface: The framework subsumes a range of fixed- and hierarchical-horizon strategies within a single trained model, eliminating the need for separate high-level and low-level policies or auxiliary modules.
Future Direction: The results suggest that variable-length supervision is a foundational step toward more robust planning algorithms, offering a flexible interface where planners can adapt prediction granularity to the specific needs of the task.