Source-linked AI summary
Predicting Consequences and Reinforcing Navigation Policies with Latent World Models
Zengmao Wang, Wei Gao, Shuhan Shen
TL;DR
Existing world models often reconstruct future observations or features, although navigation requires evaluating counterfactual action consequences. LWM predicts action-conditioned latent compatibility using cross-trajectory imagination, then supports policy learning from unlabeled videos and reinforcement within the model. The paper reports significant improvements over prior methods in prediction accuracy, policy learning, and real-world robotic navigation.
Problem
Existing world models reconstruct future observations from factual dataset actions, creating a mismatch with counterfactual planning and decision-making.
Method
LWM predicts action-conditioned latent feature compatibility using cross-trajectory counterfactual action sequences instead of reconstructing future observations.
Results
LWM significantly outperforms existing world models in real-world robotic navigation, prediction accuracy, and policy learning.
Takeaways & Limitations
The learned model supports policy supervision from unlabeled videos and imagination-driven reinforcement learning without additional environment interaction.
Takeaways & Limitations
LWM assumes spatial proximity strongly correlates with visual similarity, which may weaken in densely cluttered, partially observable environments with severe occlusions.
Abstract
from arXiv · showhide
World models enable agents to reason about future outcomes and learn policies from their knowledge of state transition, but existing approaches primarily focus on reconstructing future observations or features, which introduces unnecessary complexity and limits their effectiveness for decision making. In this work, we propose a compatibility prediction Latent World Model (LWM) for robot navigation that predicts action-conditioned latent feature compatibility rather than reconstructing observations. Our key insight is that spatial proximity correlates with latent feature similarity, enabling action consequences to be evaluated directly in latent space. To support counterfactual training, our model leverages action sequences sampled across trajectories and learns to predict which sequences lead closer to the goal. Furthermore, we demonstrate how the learned world model can supervise policy learning from unlabeled video data and further improve policies through reinforcement learning entirely within the world model. This imagination-driven framework eliminates the need for action annotations and additional environment interaction. Extensive experiments on multiple real-world robot navigation datasets show that our approach significantly outperforms prior world model and imitation learning methods in prediction accuracy, policy learning, and real-world navigation performance. The code, pretrained models, and additional materials are available at https://wzm206.github.io/latent-world-model-nav.
1 Introduction
The paper proposes a Latent World Model that predicts action-conditioned latent compatibility for counterfactual navigation reasoning instead of reconstructing future observations. It also uses the model to supervise and reinforce policies from offline, unlabeled data, reporting improvements in prediction, policy learning, and real-world navigation.
- Motivation: Existing world models reconstruct factual future observations, whereas planning requires evaluating counterfactual actions and decision-oriented consequences.Reconstruction objectives also introduce complexity, compounding errors, and representations misaligned with decision-making.
- Core formulation: LWM predicts compatibility between action consequences and future latent states, motivated by the correlation between spatial proximity and visual feature similarity.This replaces exact observation generation with direct latent-space compatibility prediction.
- Counterfactual training: Cross-trajectory action sequences let LWM simulate alternative futures and learn which candidates lead closer to the ground-truth future from offline data.The framework introduces counterfactual reasoning during training without requiring additional real-world interaction data.
- Policy learning: The learned world model labels candidate action sequences for supervised policy learning from unlabeled videos and supports subsequent reinforcement learning entirely within the model.The model predicts similarity to goal states and serves as a reward model for imagination-driven policy improvement.
- Evaluation: LWM achieves state-of-the-art performance in real-world environment and robot navigation benchmarks.The evaluation covers real environments and three robot navigation datasets.
2 Related Work
Prior world-model research spans representation learning, predictive modeling, video generation, and policy optimization using imagined or synthesized futures. LWM differs by avoiding pixel-space generation and feature reconstruction, enabling rapid latent-space prediction and decision-making.
- World-model scope: World models learn internal environmental representations and use predictive modeling to guide agents’ external decisions.This broad framing includes both understanding the environment and supporting action selection.
- Generative approaches: Video-generation world models condition future-frame prediction on observations and actions, but pixel-level generation can be slow and physically inconsistent.Feature-generation methods such as Dino-WM reduce pixel-space demands while retaining reconstruction as their primary objective.
- LWM distinction: LWM avoids pixel-space generation and feature reconstruction, eliminating iterative diffusion for rapid prediction and decision-making.Its compatibility-prediction objective focuses on action consequences rather than reconstructing future features.
- Policy optimization: Model-based reinforcement-learning methods learn dynamics and optimize policies through imagined trajectories, supporting long-horizon policy improvement.Examples include PlaNet, Dreamer, DreamerV2, and PWM.
- World models as simulators: Frozen world models can act as simulators for policy training, often using synthesized pixel-level information for control or sim-to-real transfer.Representative systems include DreamZero, World-env, RWM, and WorldVLA.
3 Method
The proposed Latent World Model predicts action-conditioned latent compatibility instead of reconstructing future observations, enabling counterfactual evaluation from offline trajectories. It supports goal-directed planning, unlabeled-video policy learning, and reinforcement learning entirely in latent space.
- Latent World Model: LWM predicts action-conditioned latent feature compatibility rather than reconstructing future observations, bringing decision-making into latent space.The framework evaluates predicted features against a goal representation instead of generating pixels or future observations.
- World Model Formulation: The world model predicts future state features for arbitrary action sequences and discriminates their compatibility with a goal observation.Its prediction module produces future features, while its discrimination module scores those features against the goal.
- Counterfactual Training: Counterfactual training samples candidate action sequences from other trajectories and learns which predicted futures are closer to the ground-truth future.This exposes the model to alternatives unavailable in each individual offline trajectory without environment interaction.
- Compatibility Learning Objective: Compatibility learning treats the ground-truth sequence as a positive and counterfactual sequences as negatives, using Info-NCE and spatial-distance-based labels.A logarithmic mapping of Euclidean distance makes supervision sensitive to small deviations and spatially close hard negatives.
- Planning: For planning, the model scores candidate action sequences against a goal and executes the sequence with the highest compatibility.The goal may be supplied as an image and can be extended to semantic targets or language instructions.
- Policy Learning: The learned model labels action sequences for supervised policy learning from unlabeled video and supports further reinforcement learning entirely within the latent world model.The policy is goal-conditioned and outputs action sequences; latent reinforcement increases the likelihood of sequences leading to goal-consistent futures.
4 Experiment
Experiments evaluate LWM across three datasets, world-model prediction, policy learning, reinforcement, real-world navigation, generalization, and ablations. LWM achieves stronger reported prediction, policy, and navigation performance while benefiting from counterfactual action sequences.
- Experimental Setup: Experiments use RECON, SCAND, and the authors’ LWM dataset, spanning outdoor, campus, residential, park, and dynamic-object scenarios.The datasets contain approximately three hours of wheeled-robot data for RECON and SCAND and approximately six hours for LWM.
- World Model Evaluation: LWM achieves the lowest prediction errors and highest action-selection accuracy among the compared world models.Evaluation reports Position Error, Orientation Error, ACC3, and ACC5 for candidate action-sequence selection.
- Policy Learning: LWM provides reliable supervisory signals for policy training from unlabeled video data and achieves the lowest reported policy errors.The policy-learning evaluation reports Position Error and Orientation Error under supervision from different world models.
- Policy Reinforcement: Performance improves across policy initializations when policies are further optimized with GRPO inside LWM.The world model is used as a unified reward function for reinforcement learning within the model.
- Real-World Navigation: In six real-world navigation scenarios, LWM achieves the highest success rate and lowest LPIPS without action annotations during training.The policy is learned from video data through the world model and remains robust when goals are far from the initial position.
- Generalization: LWM reaches the same indoor goals from different starting positions despite those environments being absent from training data.This result is presented as evidence of learned spatial and geometric relationships rather than memorized action sequences.
- Ablations: Using multiple imagined action sequences improves performance by providing richer comparative supervision for counterfactual reasoning.The ablation defines M > 1 as using counterfactual action sequences during training.
- Ablations: The reported loss-function ablation finds contrastive loss has higher prediction error than MSE, while ranking loss improves relative ordering among action sequences.The table caption identifies MSE with log as the best-performing proposed method for capturing spatial relationships and subtle action consequences.
5 Conclusion
The paper presents LWM as a compatibility-prediction world model for visual navigation. Its cross-trajectory counterfactual reasoning supports unlabeled-video policy supervision and imagination-based reinforcement learning, with reported gains in real-world navigation.
- 5 Conclusion: LWM predicts action-conditioned latent feature compatibility through cross-trajectory counterfactual reasoning instead of reconstructing observations.The framework supports policy supervision from unlabeled videos and reinforcement learning within the learned world model.
A Details of the Model
The LWM framework contains a world model and a policy model, with variants differing in visual encoder capacity.
- A Details of the Model: The framework comprises two primary components: the world model and the policy model.The section states that architectural specifications and training hyperparameters are provided for both components.
- A Details of the Model: LWM-S uses DINOv2-small with 49M parameters, while LWM-B and LWM-B-P use CroCo-v2-base with 217M parameters.The encoder choices are used to compare models with different capacities and to match the Dino-WM parameter count for LWM-S.
A.2 Policy Model
The policy comparison uses diffusion, behavioral-cloning, and LWM policy architectures with a common visual encoder. LWM autoregressively generates action tokens with a Transformer decoder.
- Policy Architectures: All three policy methods use the CroCo-v2-base visual encoder for fair comparison.The evaluated architectures are NoMaD, Behavioral Cloning, and LWM Policy.
- Policy Architectures: NoMaD combines encoded current and goal images with a diffusion-policy Transformer and one-dimensional U-Net.The U-Net estimates noise from noisy actions and the diffusion timestep.
- Policy Architectures: Behavioral Cloning encodes current and goal observations, pools the features, and projects them into action sequences with an MLP.Its Transformer encoder has 8 layers, 6 attention heads, and feature dimension 384.
- Policy Architectures: LWM Policy uses a six-layer Transformer decoder to autoregressively generate action tokens from encoded image features.Generation begins with [START] and ends when [END] is produced.
- Training: All policy models are trained for 40 epochs with Adam, batch size 24, and a 1 × 10^-4 learning rate with four-epoch warmup.The training schedule is shared across the policy models.
- Action Representation: The action-sequence centers used in the experiments are obtained through k-means clustering.
B Ablation of counterfactual action training
Counterfactual action-sequence training is evaluated through ablations on LWM-S and LWM-B-P. Removing this training severely degrades performance, while adding more counterfactual sequences improves results before saturation.
- Ablation experiments cover counterfactual action sequences in both LWM-S and LWM-B-P.
- Removing counterfactual training severely degrades model performance.
- Adding counterfactual training yields progressive performance gains that eventually saturate.The results suggest that compatibility labels provide strong initialization rather than an absolute upper bound.
C Navigation in real world
On the LWM test dataset, the world model predicts a smooth compatibility surface whose peak lies near the ground-truth value. Lower distance-function values represent higher compatibility, supporting policy learning.
- The visualization uses a distance function in which lower z-axis values indicate higher compatibility.
- The predicted compatibility surface is smooth across the real-world test dataset.
- Maximum compatibility peaks near the ground-truth value, a property described as advantageous for policy learning.
D Limitations and Future Work
The method depends on spatial proximity correlating with visual similarity, but this relationship can weaken under severe occlusion and dense clutter. The paper suggests 3D-aware representations or long-term temporal memory as possible remedies.
- The approach assumes that spatial proximity strongly correlates with visual similarity.
- Dense clutter, severe occlusions, and marginal spatial displacements can cause abrupt visual changes that reduce compatibility discrimination.
- Future work could integrate 3D-aware latent representations or long-term temporal memory for occlusion-induced visual shifts.