Unified World Models: Memory-Augmented Planning and Foresight for Visual Navigation

TL;DR

This work tackles state-action misalignment in visual navigation by introducing UniWM, a unified memory-augmented world model that grounds planning in visually imagined futures within a single multimodal backbone. It employs a hierarchical memory to fuse immediate perceptual cues with long-horizon trajectory context and trains via interleaved planner and world-model objectives with discretized action tokens and image reconstruction. Empirical results across four benchmarks show significant navigation and visualization gains, plus zero-shot generalization to unseen environments like TartanDrive, underscoring the method's robustness and scalability. Ablation analyses confirm the value of memory, the planning and visualization losses, and the interleaved training scheme as key drivers of performance.

Abstract

Enabling embodied agents to effectively imagine future states is critical for robust and generalizable visual navigation. Current state-of-the-art approaches, however, adopt modular architectures that separate navigation planning from visual world modeling, leading to state-action misalignment and limited adaptability in novel or dynamic scenarios. To overcome this fundamental limitation, we propose UniWM, a unified, memory-augmented world model integrating egocentric visual foresight and planning within a single multimodal autoregressive backbone. Unlike modular frameworks, UniWM explicitly grounds action decisions in visually imagined outcomes, ensuring tight alignment between prediction and control. A hierarchical memory mechanism further integrates detailed short-term perceptual cues with longer-term trajectory context, enabling stable, coherent reasoning over extended horizons. Extensive experiments across four challenging benchmarks (Go Stanford, ReCon, SCAND, HuRoN) demonstrate that UniWM substantially improves navigation success rates by up to 30%, significantly reduces trajectory errors compared to strong baselines, and exhibits impressive zero-shot generalization on the unseen TartanDrive dataset. These results highlight UniWM as a principled step toward unified, imagination-driven embodied navigation.

Figures (11)

• Figure 1: Comparison of goal-conditioned visual navigation methods. All panels use the same start/goal observations; headers report navigation performance SR$\uparrow$, ATE$\downarrow$, and RPE$\downarrow$ on HuRoN hirose2023sacson dataset. (a) Navigation policy methods like NoMaD sridhar2024nomad directly predict action sequences $A_T$. (b) World model for navigation like NWM bar2025navigation uses a world model to visualize future observations, enhancing a separate navigation planner. (c) UniWM (no memory) unifies planning and visualization within one multimodal backbone, and actions are grounded in the imagined next observation while generating $A_T$ autoregressively. (d) UniWM (with hierarchical memory) adds intra-step and cross-step memory banks, stabilizing longer-horizon rollouts and consistently yielding the highest SR and lowest errors (ATE/RPE).
• Figure 2: UniWM framework.(a) Training: planner and world-model samples are interleaved within a single unified multimodal autoregressive backbone, optimized jointly with the discretized bin-token loss $\mathcal{L}_{\text{plan}}$ and the reconstruction loss $\mathcal{L}_{\text{world}}$; bin/text/image tokenizers map actions, pose, and observations to tokens. (b) Inference: a hierarchical memory supplies intra- and cross-step KV states ($\mathcal{M}^{\text{intra}}_{t}$ caches the current observation; $\mathcal{M}^{\text{cross}}_{t}$ accumulates prior steps) to augment attention, yielding robust trajectory-consistent alternating predictions of $\hat{a}_t$ (next action) and $\hat{o}_t$ (next observation). See Fig. \ref{['fig:detail']} for the detailed memory mechanism.
• Figure 3: Overview of hierarchical memory bank mechanism ($\mathcal{M}^{\text{intra}}_t$ & $\mathcal{M}^{\text{cross}}_t$).(a)$KV$ (keys/values) extracted from selected layers are deposited into $\mathcal{M}^{\text{intra}}_t$ at the beginning of each step $t$ (Eq. \ref{['eq:seg_kv']}). (b)(c)$\mathcal{M}^{\text{intra}}_t$ is merged with the accumulated cross-step memory $\mathcal{M}^{\text{cross}}_t$ via top-$k$ similarity gating (Eq. \ref{['eq:sim_gate_topk']}) and exponential temporal decay (Eq. \ref{['eq:decay']}), yielding a fused memory (Eq. \ref{['eq:fuse']}) that augments attention for both the planner and the world-model substeps (Eq. \ref{['eq:enhance_atten']}) to promote trajectory-consistent predictions. At the end of step $t$, $\mathcal{M}^{\text{intra}}_t$ (with timestamp $t$) is appended to $\mathcal{M}^{\text{cross}}_t$ for reliable reuse at step $t{+}1$, enabling robustly efficient rollouts.
• Figure 4: Qualitative Comparisons on Go Stanford and HuRoN across UniWM, NWM, and NoMaD. The central trajectory plots highlight difference between predicted $A_T$ and the ground-truth.
• Figure 5: Impact of discretized bin-token loss ($\mathcal{L}_{\text{plan}}$) and reconstruction Loss ($\mathcal{L}_{\text{world}}$) on navigation (left) and visualization (right) performance, averaged over evaluation splits of Go Stanford, ReCon, SCAND, and HuRoN. X-axis arrows indicate whether higher or lower values are preferable.