NestedWM: World Models with Nested Latent Representations

Planning with world models requires repeated rollouts, yet conventional models are trained with a fixed latent size that determines the trade-off between representation capacity and planning cost. We propose NestedWM, which learns nested latent representations that allow a single world model to predict and plan using leading prefixes of its predictive state at different budgets. Each prefix contains either the leading dimensions of a latent vector or the leading tokens of an ordered latent-token sequence. We jointly train an encoder and an action-conditioned predictor over sampled prefix budgets, using each selected prefix as both the predictor input and the target prediction space. In a dimension-wise study with the architecture held fixed, a single encoder and predictor support planning across budgets, achieving high success with only a fraction of the full latent dimensions in most evaluated visual-control and synthetic tasks. Further analyses show that leading prefixes preferentially retain planning-relevant information, while predictable distractors can consume limited representation budget. We then turn this capability into computational savings with a token-wise instantiation that processes only the selected prefix tokens during rollout. In controlled multi-token settings, a single model supports successful planning across multiple token budgets, with smaller prefixes reducing model-based cost-evaluation time.