Space-Aware World Models: Spatial Persistence Through Factored Scene Representations
Abstract
Long-horizon navigation simulations require spatial persistence; revisited physical states must reproduce the same landmarks, rather than merely locally plausible frames. Standard latent-dynamics world models typically evolve a single predictive state that carries the scene layout, appearance, and agent motion simultaneously, which makes revisitation consistency brittle under compounding errors. We study a structured regime for navigation-like tasks in which the scene layout remains static during an episode, a map-like scene observation is available, and the agent motion can be described by a low-dimensional state. For this regime, we propose Space-Aware World Models (SWM). SWM factors world modeling into: (i) a scene latent inferred once per episode from a top-down map; (ii) an explicit, action-conditioned transition model over agent state; and (iii) a conditional renderer from scene code and state to observations. By anchoring scene content in an episode-level representation, SWM is designed to preserve the identity of revisited locations across autoregressive imagined trajectories. Across procedurally generated environments, SWM improves long-horizon fidelity and yields stronger revisitation and geometric consistency, as well as control-relevant fidelity, outperforming standard latent-dynamics baselines by eliminating compounding scene-generation errors. These results suggest that separating the scene from state evolution provides an effective inductive bias for spatially coherent simulations.