Source-linked AI summary
Looped World Models
Hongyuan Adam Lu, Z. L. Victor Wei, Qun Zhang, Jinrui Zeng, Bowen Cao, Lingwei Meng, Mocheng Li, Zezhong Wang, Haonan Yin, Naifu Xue, Minyu Chen, Cenyuan Zhang, Zefan Zhang, Hao Wei, Jiawei Zhou, Haoran Xu, Hao Yang, Ronglai Zuo, Tongda Xu, Yonghao Li, Jian Chen, Hebin Wang, Zeyu Gao, Yang Li, Wei Zhao, Qimin Zhong, Siqi Liu, Yumeng Zhang, Leyan Cui, Zhangyu Wang, Wai Lam
TL;DR
Long-horizon world simulation requires substantial computation, while existing looped architectures have not been explored for world modelling. LoopWM iteratively refines latent environment states with a parameter-shared transformer and achieves up to 100× parameter efficiency over conventional approaches.
Problem
Existing world models face unresolved tension between computational depth for faithful long-horizon simulation and rollout stability, deployment cost, and unexplored looped architectures for world modelling.
Method
LoopWM iteratively refines latent environment states through a parameter-shared transformer block with stabilized residual dynamics and adaptive computation.
Results
Up to 100× parameter efficiency is achieved over conventional approaches without sacrificing prediction quality, alongside stable long-horizon rollouts and adaptive computation.
Takeaways & Limitations
Iterative latent depth emerges as a scaling axis for world simulation that is orthogonal to model size and data volume.
Takeaways & Limitations
The paper has selective presentation scope, with broader cross-family positioning, richer scaling analysis, and more extensive optimization disclosure left for future work.
Abstract
from arXiv · showhide
Current world models face a fundamental tension: faithful long-horizon simulation demands deep computation, but deeper models are expensive to deploy and prone to compounding errors. We resolve this by introducing Looped World Models (LoopWM), which are the first looped architectures for world modelling. Our method iteratively refines latent environment states through a parameter-shared transformer block. This yield up to 100x parameter efficiency over conventional approaches with adaptive computation that automatically scales depth to match the complexity of each prediction step. Orthogonal to scaling model size and training data, LoopWM establishes iterative latent depth as a new scaling axis for world simulation, which might significantly push the community forward.
1 INTRODUCTION
LoopWM introduces looped transformer architectures for world models, addressing the need for iterative computation in faithful long-horizon simulation. It combines parameter sharing, iterative latent refinement, and adaptive depth for environment dynamics prediction.
- Motivation: Faithful long-horizon simulation often requires deep or iterative computation because fixed-depth architectures use identical computation for transitions of different difficulty.Extended rollouts also suffer rapid trajectory-quality degradation from compounding prediction errors.
- Motivation: Previous looped architectures were developed and evaluated exclusively for language modelling, leaving Looped World Models entirely unexplored.This motivates applying recurrently shared transformer blocks to environment simulation and dynamics prediction.
- Method: LoopWM models a single-step transition through repeated application of a shared latent update operator, explicitly refining environment states while reusing parameters across depth.The correspondence between environment dynamics and the inner loop is conceptual rather than exact.
- Motivation: Hundreds or thousands of sequential dynamics-model executions make parameter efficiency especially valuable for long-horizon world-model rollouts.Adaptive depth can allocate more iterations to complex transitions such as collisions.
- Contribution: LoopWM combines a parameter-shared recurrent transformer block with spectrally constrained residual dynamics for provably stable state transitions across arbitrary rollout lengths.The paper reports competitive or superior predictive accuracy while using significantly fewer resources than existing architectures.
2 RELATED WORK
Related work spans latent, recurrent, transformer-based, and generative world models, while highlighting compounding rollout error as a persistent challenge. Looped transformers provide a parameter-sharing foundation for reducing parameter count and adapting computation to input complexity.
- World Models: World-model research progressed from cognitive-science forward models to latent-space reinforcement-learning systems such as PlaNet, SimPLe, and MuZero.Ha and Schmidhuber learned compressed environment representations for policy training inside a learned dream; PlaNet planned directly in latent space, and SimPLe trained a video-prediction simulator.
- World Models: Dreamer refined RSSM-based world models, while later systems introduced autoregressive and state-space transformers.DreamerV3 achieved human-level performance across over 150 diverse tasks with one hyperparameter set; IRIS replaced recurrence with an autoregressive transformer over discrete tokens.
- World Models: Video-generation models such as Sora, Genie, and Genie 3 extended world modelling toward long-form, interactive, and explorable environments.Sora demonstrated emergent 3D consistency, while Genie and Genie 3 produced text-conditioned interactive environments.
- Open Challenges: Compounding prediction error remains a persistent challenge because small rollout inaccuracies accumulate exponentially over long horizons and degrade trajectory fidelity.Mitigation strategies include short-horizon re-planning, self-correcting models, and physics-informed architectures, but the passage identifies a continuing computational tension.
- Looped Transformers: Looped transformers reuse shared blocks across depth, decoupling effective computation from parameter count and enabling input-dependent halting through Adaptive Computation Time.The Universal Transformer introduced weight sharing with ACT, while ALBERT demonstrated full cross-layer sharing in BERT-scale models; later analyses showed looped transformers can simulate arbitrary programs and match in-context-learning performance with less than 10% of the parameters.
- Adaptive Computation: Adaptive computation has also been studied through per-step halting in RNNs and early-exit methods that terminate inference at intermediate layers when confidence is sufficient.These approaches allocate variable computation to inputs with differing complexity.
3 LOOPED WORLD MODEL
Looped World Models combine iterative looped-transformer computation with action-conditioned state prediction for world modelling. Their design aligns computation with physical dynamics, targets stable latent transitions over arbitrary rollout lengths, and adapts computational depth to prediction complexity.
- Architecture: Looped World Models combine iterative looped-transformer computation with action-conditioned state prediction for world modelling.The architecture is presented as a latent dynamics model.
- Design principles: The design aligns the model’s computation graph with the iterative nature of physical dynamics.This is the first stated design principle.
- Design principles: The model targets provable stability of latent state transitions across arbitrary rollout lengths.This is the second stated design principle.
- Design principles: Adaptive computational depth matches the complexity of each prediction step.This is the third stated design principle, though the supplied passage truncates before specifying the matching mechanism.
3.1 OVERALL ARCHITECTURE
LoopWM predicts future latent states from observations and actions using four modules. Its central looped dynamics core iteratively updates the latent state, which prediction heads decode into observations, rewards, and continuation signals.
- Overall architecture: The architecture comprises an observation encoder, action embedder, looped dynamics core, and prediction heads.The encoder maps o_k to e_k, while the action embedder maps a_k to u_k in the shared latent space R^d.
- Overall architecture: The looped dynamics core produces h_k from h_{k−1}, e_k, and u_k through T iterations of a parameter-shared transformer block.Its residual dynamics are spectrally constrained.
- Overall architecture: The prediction heads decode h_k into a reconstructed observation or latent target, predicted reward, and predicted continuation flag.These lightweight MLP heads follow standard latent world-model designs.
- Overall architecture: At environment step k, the forward pass computes e_k = Eϕ(o_k), u_k = Aψ(a_k), h_k = Lθ(h_{k−1}, e_k, u_k), then decodes the predictions.The complete mapping is summarized by Equation (3).
3.2 LOOPED DYNAMICS CORE WITH SPECTRAL STABILITY
The dynamics core uses a three-block prelude–recurrent–coda design, with a parameter-shared recurrent transformer refining latent states over T loops. Spectral constraints keep recurrent dynamics bounded and support state propagation across environment timesteps.
- Architecture: The dynamics core is partitioned into Prelude P, Recurrent Block R, and Coda C.The prelude, recurrent, and coda blocks respectively condition, iteratively refine, and project the latent state.
- Prelude P: Prelude P processes the previous latent state, observation embedding, and action embedding to produce the conditioning signal e.Layer normalization of e prevents input magnitude from inducing late-stage loss spikes.
- Recurrent Block R: Recurrent Block R applies shared transformer parameters iteratively for T loops, with state retention and input injection controlling each update.The hidden state is initialized randomly or from the previous timestep’s final hidden state, while shared parameters make computational depth independent of parameter count.
- Spectral Stability Constraint: The spectral norm constraint on Ā is strictly less than 1, ensuring bounded latent dynamics for any loop count T.Ā is constructed from a negative diagonal matrix, giving diagonal entries in (0, 1) and ρ(Ā) < 1 without gradient clipping or post-hoc normalization.
- Cross-Timestep State Propagation: Terminal state h(T) can initialize the next environment step, creating inner-loop refinement and outer-loop temporal propagation that remain bounded under the spectral constraint.The coda projects h(T) for prediction, and the resulting state is carried forward across time.
3.3 TRAINING OBJECTIVE
LoopWM trains with stochastic, sequence-level loop depths and a multi-task world-model objective. Truncated backpropagation controls memory, while entropy regularization stabilizes adaptive early exit.
- Variable-Depth Training: Loop depth T is sampled independently per sequence from a Poisson distribution with learnable mean µrec.Sequence-level sampling reduces training-objective variance and empirically eliminates most loss spikes.
- World Model Loss: The world-model loss combines observation, reward, and continuation prediction.The observation-loss form depends on the observation space, such as MSE for continuous states or cross-entropy for discrete tokens.
- Entropy-Regularised Adaptive Depth: Entropy regularization prevents adaptive exit gates from collapsing to always exiting at the first iteration or never exiting.With adaptive early exit enabled, the total loss is L = Lwm + Lent.
3.4 ADAPTIVE EARLY EXIT FOR INFERENCE
LoopWM adaptively allocates inference iterations by terminating quickly converging transitions and extending computation for complex ones. This enables substantial FLOPs savings while supporting progressively refined predictions through test-time compute scaling.
- Adaptive early exit: A lightweight exit gate adaptively terminates the inner loop when its sigmoid output g(t) exceeds threshold τ, using h(t) as the final hidden state.The gate is a single-layer MLP followed by a sigmoid, with learned parameters w_g ∈ R^d and b_g ∈ R.
- Adaptive early exit: The exit mechanism complements convergence-based criteria that halt when the second-order difference ∥h(t) − 2h(t−1) + h(t−2)∥ falls below a threshold.This provides an alternative stopping signal based on latent-state dynamics.
- Inference savings: Up to two orders of magnitude aggregate FLOPs reduction is possible over long rollouts containing many simple transitions and occasional complex events, compared to a fixed-depth model.Adaptive computation concentrates additional iterations on complex transitions while exiting simple ones early.
- Test-time compute scaling: T_max can exceed the training-time mean μ_rec at inference, enabling test-time compute scaling with progressively refined predictions as more iterations are allocated.The maximum loop count is not restricted to the average number of training-time recurrent iterations.
3.5 DEFERRED DECODING: ACTION-CONDITIONED LATENT ROLLOUT
Deferred Decoding removes intermediate observation, reward, and continuation decoding during multi-step action-conditioned rollouts, allowing uninterrupted latent dynamics and a single terminal decode. Its nested computation combines K action steps with T inner refinements while training adds trajectory regularization and a progressive horizon curriculum.
- Motivation: The approach addresses per-step decoding inefficiencies by freeing intermediate latent states from pixel reconstruction and enabling uninterrupted reasoning across action sequences.Standard decoding forces latent capacity toward intermediate pixel-level reconstruction even when planning uses only the final prediction.
- Deferred decoding: Deferred Decoding advances action-conditioned latent states across K steps without intermediate outputs, then decodes observation, reward, and continuation predictions only at step K.The decoder is invoked exactly once during the rollout.
- Nested loops: The nested rollout has K outer action steps and T inner latent refinements, yielding K × T shared-parameter transformer applications with one decoder forward pass.Each outer step injects a new action embedding before the recurrent latent-refinement loop.
- Training: Training combines terminal prediction with latent trajectory constraints to preserve semantic consistency and bounded evolution across unsupervised intermediate states.The constraints include latent consistency guidance and a cumulative contraction penalty over deferred horizons.
- Training: Training begins at K = 1 and progressively increases the horizon according to K(step) = min(Kmax, 1 + ⌊step/∆⌋) to stabilize gradients through K × T applications.The curriculum first teaches accurate single-step transitions before longer deferred rollouts.
- Inference modes: In planning mode, decoder invocations fall from K to 1, saving approximately (K −1)×cost(Dξ) FLOPs per candidate sequence.With adaptive early exit inside inner loops, total FLOP reduction can reach up to two orders of magnitude for long-horizon planning with simple transitions.
4 RESULTS
LoopWM outperforms claude-opus-4-6-max on ScienceWorld, including a 21.2% average EM advantage and a Lifespan increase from 0% to 100% in extreme cases. On AlfWorld, the roughly 1B-parameter model achieves the best BLEU scores among four models and ranks second on EM and Token F1, while deferred decoding becomes more useful with accumulated rollouts.
- ScienceWorld: 21.2%: LoopWM surpasses claude-opus-4-6-max on average EM on ScienceWorld.The passage also reports improvements on other metrics.
- ScienceWorld: 0% to 100%: LoopWM improves ScienceWorld Lifespan scores in the most extreme cases.The authors present this as evidence of the model’s underlying capacity.
- ScienceWorld: LoopWM remains competitive across metrics and is much stronger than qwen-3.5-flash and gemini-3-flash-preview.The passage attributes the weaker baseline performance to their smaller model sizes.
- AlfWorld: Around 1B parameters: the proposed AlfWorld model achieves the best BLEU metrics among four models and ranks second on EM and Token F1.Detailed action analysis identifies low entity scores as a main area for future optimization.
- AlfWorld: EM: 68.4%; Token F1: 85.3%; BLEU-4: 80.7%; Entity: 83.9%.These are the reported overall results for one evaluated model.
- Decoding: Deferred decoding is useful, particularly when rollouts are accumulated.The conclusion is drawn across the reported tables.
5 CONCLUSIONS
LoopWM applies looped transformer architectures to world modelling by iteratively refining latent environment states with parameter-shared computation. It adaptively matches refinement depth to prediction complexity and proposes iterative latent depth as a scaling axis orthogonal to model size and data volume.
- 5 CONCLUSIONS: LoopWM is the first application of looped transformer architectures to world modelling, iteratively refining latent environment states through a parameter-shared transformer block.The approach uses stabilised residual dynamics to address the tension between faithful long-horizon simulation, deployment cost, and compounding rollout errors.
- 5 CONCLUSIONS: LoopWM automatically scales effective depth to prediction complexity, allocating more refinement iterations to challenging transitions and fewer to predictable ones.This mechanism adjusts computation dynamically across prediction steps.
- 5 CONCLUSIONS: The work identifies iterative latent depth as a new scaling axis for world simulation, orthogonal to model size and data volume.The authors present this perspective as a direction for building future world models.
6 BROADER IMPACTS
The section argues that LoopWM’s core architectural thesis is already supported, while broader disclosure of evidence, cross-family positioning, scaling laws, and optimization results remains future work.
- Scope of current evidence: LoopWM’s core case is established through looped latent refinement, deferred decoding, and stabilized dynamics, but the manuscript selectively discloses supporting results.The authors frame the limitation as presentation scope rather than a lack of conceptual or empirical foundation.
- Iterative latent computation: Deferred decoding provides concrete evidence that preserving and refining latent computation across rollout steps benefits world modelling.The authors describe deferred decoding as one visible entry point into broader evidence supporting looping.
- Evidence beyond reported domains: Optimization is feasible in continuous visual environments, with consistently reducible training loss beyond the principal task domains reported.The remaining limitation is described as incomplete disclosure rather than insufficient empirical support.
- Positioning in the world-model landscape: A fuller analysis should position LoopWM against RSSM-style, autoregressive video-token, and diffusion-based world models and clarify where iterative latent depth is the natural scaling axis.The paper identifies broader cross-family positioning and regime clarification as opportunities for future work.
- Scaling and optimization: Step 1 to step 5 experiments indicate that iterative latent depth is a meaningful scaling dimension, but broader task and compute ranges are needed for complete scaling laws.The authors also identify more extensive optimization disclosure as remaining work.