🤖 AI

Latent Goal Prediction from Language for Model-Based Planning

The paper introduces LAGO, a framework that enables robust long-horizon model-based planning by dynamically decomposing language instructions into intermediate latent subgoals and action-conditioned rollouts within a unified latent space, thereby overcoming the limitations of both visual targets and noisy cross-modal alignment.

Original authors: Samuel Barbeau, Simon Roy, Giovanni Beltrame, Christian Desrosiers, Nicolas Thome

Published 2026-06-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Samuel Barbeau, Simon Roy, Giovanni Beltrame, Christian Desrosiers, Nicolas Thome

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 to navigate a complex maze or move objects around a room, but you can only give it instructions in plain English, like "Go to the red dot" or "Put the cube in the drawer."

This paper introduces a new system called LAGO (Latent Goal Prediction from Language) that helps robots plan these long, complicated journeys much better than before. Here is how it works, broken down into simple concepts:

The Problem: The "Lost in Translation" and "Long Walk" Issues

Currently, robots trying to plan ahead face two main headaches:

The "Too Far" Problem: If you ask a robot to go somewhere very far away, it tries to imagine the whole trip in one giant leap. But just like trying to predict the weather a month from now, small mistakes in its imagination pile up, and the plan falls apart before it even starts.
The "Language vs. Picture" Problem:
- If you give the robot a picture of the goal, it knows exactly where to go, but it's rigid. It can't handle new situations easily.
- If you give the robot words, it's flexible, but it's hard to translate "Go to the dragon" into precise coordinates for the robot's brain. Existing methods often get confused or require massive, slow computers to understand the words.

The Solution: LAGO's "Stop-and-Check" Strategy

LAGO solves this by acting like a hiker with a map and a series of checkpoints.

Instead of trying to imagine the entire 100-mile hike in one go, LAGO breaks the journey down into small, manageable steps.

The "Latent" Map: The robot doesn't think in raw pixels (like a camera sees) or raw words. It thinks in a "secret code" (a latent space) that represents the world.
The Translator: LAGO is trained to take your English sentence (e.g., "Go to the villager") and instantly translate it into a sequence of invisible checkpoints in that secret code.
The Process:
1. You say, "Go to the villager."
2. LAGO doesn't just say "Okay." It predicts a path: "First, imagine being here (Checkpoint 1), then here (Checkpoint 2), then here (Checkpoint 3)..." all the way to the villager.
3. The robot plans its next move to reach the first checkpoint.
4. Once it gets there, it updates its position and predicts the next set of checkpoints based on where it actually is now.

The Analogy: The "Smooth Valley" vs. The "Rough Mountain"

The paper uses a great visual analogy to explain why this works better:

Old Methods: Imagine trying to roll a ball to the top of a mountain by guessing the whole path at once. The terrain is bumpy and full of holes (errors). The ball gets stuck or rolls the wrong way because the path is too long and complex to see clearly.
LAGO: Imagine the same mountain, but LAGO digs a smooth, winding valley with clear stepping stones (the subgoals) leading to the top. The robot just needs to hop from one stone to the next. Even if it misses a stone slightly, the valley guides it back on track. It never has to look at the whole mountain at once; it only looks at the next stone.

Why This is a Big Deal

It's Fast: Unlike other systems that use huge, slow AI models to understand language, LAGO is lightweight and fast, making it suitable for real-time planning.
It's Robust: In tests, when the distance to the goal got very long, other methods failed completely (0% success). LAGO kept succeeding, even when the journey was difficult.
It Bridges the Gap: It combines the best of both worlds: the flexibility of understanding human language and the precision of visual targets.

What It Does (and Doesn't) Do

It Does: It takes a language instruction and a current view of the world, then generates a smooth, step-by-step plan in the robot's "mind" to reach the goal. It works in simulated environments like navigating a 2D map, moving a knight in a game-like world, or pushing a cube with a robotic arm.
It Doesn't: The paper does not claim this works on real-world physical robots yet, nor does it claim to solve problems outside of these specific simulated tasks. It is a planning framework, not a physical robot itself.

In short, LAGO teaches a robot to stop trying to "see" the whole future at once and instead focus on a series of small, clear milestones predicted from your words, making long journeys much less likely to go wrong.

Technical Summary: Latent Goal Prediction from Language for Model-Based Planning

Problem Statement

Model-based reinforcement learning (RL) utilizing world models (WMs) offers a promising paradigm for sample-efficient control by simulating future trajectories in a compact latent space. However, these systems face two fundamental bottlenecks:

Goal Specification: Defining optimizable goals is difficult. Visual targets (images) provide precise local gradients but offer poor guidance for distant goals and are impractical to obtain for complex tasks. Language offers flexibility but suffers from noisy cross-modal alignment in contrastive models or computational inefficiency in large generative models, which are ill-suited for the high-frequency sampling required by model-based planning.
Long-Horizon Planning: Planning over long horizons is hindered by compounding prediction errors. Small inaccuracies in latent dynamics accumulate during rollouts, degrading the reliability of imagined trajectories. Existing methods often sidestep this by increasing temporal abstraction (frame skips) or restricting evaluation to short horizons, limiting their applicability to complex, temporally extended tasks.

Current approaches struggle to combine the precision of image-based objectives with the flexibility of language-based specification while maintaining robustness against prediction errors over long horizons.

Methodology: LAGO

The authors introduce LAGO (Latent Goal Prediction from Language), a framework designed to predict sequences of intermediate goal states from language instructions and action-conditioned rollouts within the same latent space.

Architecture

LAGO is built upon a Joint Embedding Predictive Architecture (JEPA) operating in a shared latent space. It supports two operating modes via a single predictor $\pi(\cdot)$ :

Next-Step Transition Prediction: Conditioned on low-level actions to predict the immediate next latent state.
Subgoal Prediction: Conditioned on language instructions and a completion scalar $\rho$ (estimating task progress) to predict intermediate latent subgoals.

Key Components:

Vision Encoder: Uses a frozen DINOv2 ViT to encode RGB observations into patch tokens.
Multi-Modal Conditioning: Language instructions are embedded and projected into a shared condition space. A completion scalar $\rho \in [0, 1]$ is also projected. A mode-specific embedding distinguishes between action conditioning and subgoal conditioning.
Unified Training: The model is trained jointly using a mean-squared error (MSE) objective. It minimizes the error between predicted and actual next states (using action conditioning) and between predicted and actual subgoal states (using language and $\rho$ conditioning).

Planning Strategy

LAGO transforms long-horizon planning into a sequence of locally tractable objectives:

Subgoal Generation: Given a current state and a language instruction, LAGO generates $K$ latent subgoals at uniformly spaced completion scalars ( $\rho_k = k/K$ ). Crucially, these subgoals are predicted independently from the current state (non-autoregressively) to avoid compounding errors.
Online Re-prediction: Subgoals are re-predicted at every Model Predictive Control (MPC) step, allowing the agent to adapt to its evolving state.
Soft Minimum Trajectory Cost (JLAGO): To evaluate candidate action sequences, LAGO employs a soft-minimum cost function. Instead of requiring a trajectory to hit a specific subgoal at a specific time, it measures the best alignment achieved between each subgoal and any timestep along the rollout. This is controlled by a temperature parameter $\tau$ , allowing for smooth convergence to subgoals rather than brittle discrete completion criteria.

Key Contributions

LAGO Framework: A world model trained to map natural language instructions directly to target states within its own latent space, bridging the gap between visual precision and linguistic flexibility.
Dynamic Subgoal Decomposition: An extension of goal prediction to generate sequences of intermediate latent subgoals. The approach utilizes continual online re-prediction and a soft minimum trajectory cost to enable robust long-horizon planning without explicit subgoal completion detection.
Distance-Based Evaluation Protocol: A systematic evaluation method that varies task difficulty based on the distance to the goal. This reveals that prior methods degrade sharply beyond their effective planning horizon, whereas LAGO remains robust.

Experimental Results

The authors evaluated LAGO across four environments: TwoRoom (2D navigation), Kingdom (2D navigation with dynamic layouts and irreversible hazards), OGBench Cube (manipulation), and OGBench Scene (complex manipulation).

Long-Horizon Robustness: In the TwoRoom, Kingdom, and OGBench Cube environments, image-goal baselines (DINO-WM, Temporal Straightening, LeWM) failed completely as task difficulty (distance) increased (e.g., dropping to 0% success rate at high difficulties). In contrast, LAGO maintained high success rates (e.g., 90% at difficulty 115 in TwoRoom).
Language vs. Visual Goals: LAGO conditioned on text goals significantly outperformed language-goal baselines (LIV, PE-Core, Robo-Dopamine). For instance, in Kingdom, LAGO achieved a 96.9% success rate compared to 52.6% for a finetuned LIV baseline.
Efficiency: LAGO demonstrated high inference efficiency (approx. 11ms per observation) compared to large generative reward models like Robo-Dopamine (approx. 451ms), making it viable for real-time MPC.
Ablation Studies: Performance remained stable for $K \ge 4$ subgoals. The soft-min temperature $\tau$ was found to be robust within a reasonable range.

Significance and Claims

The paper claims that LAGO successfully bridges the gap between the precision of image-based goals and the flexibility of text-guided control. By grounding language in the latent space of the world model and decomposing tasks into intermediate subgoals, the framework mitigates compounding prediction errors.

The authors state that their work enables robust and precise long-horizon planning purely from language, avoiding the sharp degradation observed in prior methods when planning horizons extend beyond immediate reach. They emphasize that this is achieved without requiring action or reward supervision during the planning phase, relying instead on the structured guidance of predicted latent subgoals. The work is presented as a step toward tackling increasingly complex planning challenges by combining the strengths of visual and linguistic modalities within a single, efficient planning framework.