Source-linked AI summary
When and How Much to Imagine: Adaptive Test-Time Scaling with World Models for Visual Spatial Reasoning
Shoubin Yu, Yue Zhang, Zun Wang, Jaehong Yoon, Huaxiu Yao, Mingyu Ding, Mohit Bansal
TL;DR
Visual spatial reasoning needs better control over when and how much to imagine because static views can be insufficient, while indiscriminate world-model use can be costly or misleading. The paper introduces AVIC for adaptive gating and planning, and AVIC-R for reinforcement-learning-based policy training; across spatial reasoning and navigation benchmarks, adaptive control achieves strong performance with fewer world-model calls and tokens, while AVIC-R can outperform proprietary policy baselines.
Problem
Visual spatial reasoning remains unreliable for unseen viewpoints, while it is unclear when imagination is necessary, how much helps, and when it becomes harmful.
Method
AVIC gates and plans world-model use from evidence sufficiency, while AVIC-R learns this policy with GRPO using QA correctness and imagination-cost signals.
Results
Across SAT, MMSI, and R2R, adaptive scaling achieves competitive or state-of-the-art performance with fewer language tokens and world-model calls, while AVIC-R outperforms proprietary policy variants in reported comparisons.
Takeaways & Limitations
World-model imagination is instance-dependent: it is most beneficial for action-conditioned reasoning, offers limited gains when observations suffice, and should be selectively controlled.
Takeaways & Limitations
The paper’s main text refers readers to the appendix for its limitations and broader-impacts discussion.
Abstract
from arXiv · showhide
Despite rapid progress in MLLMs, visual spatial reasoning remains unreliable when correct answers depend on how a scene would appear under unseen or alternative viewpoints. Recent work addresses this by augmenting reasoning with world models for visual imagination, but questions such as when imagination is actually necessary, how much of it is beneficial, and when it becomes harmful, remain poorly understood. In practice, indiscriminate imagination can increase computation and even degrade performance by introducing misleading evidence. In this work, we present an in-depth analysis of test-time visual imagination as a controllable resource for spatial reasoning. We first study when static visual evidence is sufficient, when imagination improves reasoning, and how excessive or unnecessary imagination affects accuracy and efficiency. To support this analysis, we then introduce AVIC, an adaptive test-time framework with world models that explicitly reasons about the sufficiency of current visual evidence before selectively invoking and scaling visual imagination. Finally, to further learn this gating and planning behavior without any annotation of when and how much to imagine, we introduce AVIC-R, which trains the policy via GRPO from QA-correctness rewards and penalties by imagination cost. Across spatial reasoning benchmarks (SAT, MMSI) and an embodied navigation benchmark (R2R), our results reveal clear scenarios where imagination is critical, marginal, or detrimental, and show that selective control can match or outperform fixed imagination strategies with substantially fewer world-model calls and language tokens. Our AVIC-R surpasses strong proprietary baselines including GPT-4o and GPT-4.1 while invoking the world model less often. Overall, our findings highlight the importance of analyzing and controlling test-time imagination for efficient and reliable spatial reasoning.
1 Introduction
Visual spatial reasoning remains difficult when answers depend on unseen viewpoints, while indiscriminate imagination can be unnecessary or harmful. AVIC and AVIC-R make imagination adaptive by gating and planning world-model use, with experiments showing improved efficiency and strong performance.
- Visual spatial reasoning remains challenging when answers depend on unseen regions, viewpoint changes, or transformations beyond a static observation.
- Fixed visual imagination can invoke world models without checking whether additional evidence is necessary or helpful.
- AVIC first assesses evidence sufficiency, then either answers directly or generates a dynamic-length action plan for informative imagined viewpoints.
- AVIC-R trains the gating policy with GRPO from QA-correctness rewards and imagination-cost penalties, without supervision specifying when imagination is required.
- Across SAT, MMSI, and R2R, adaptive scaling achieves competitive or state-of-the-art performance with substantially fewer language tokens and world-model calls than fixed strategies.
2 Related Work
Prior work uses world models and test-time scaling to generate novel views and improve reasoning, but typically allocates computation uniformly. This paper frames visual imagination as an adaptive resource for targeted efficiency.
- MLLMs have advanced visual tasks, yet robust spatial reasoning remains difficult and important for embodied applications.
- World models can generate future frames or action-conditioned simulations, including novel viewpoints for visual spatial reasoning.
- Existing visual test-time scaling methods commonly generate novel views and ensemble results while applying uniform computation across instances.
- Figure 2 links non-monotonic accuracy gains and higher computation from fixed imagination strategies to the need for adaptive scaling.
- Adaptive visual test-time scaling targets imagination only when necessary, improving computational efficiency.
3 Analysis of Always-on World Model Calling
Always-on world-model calling is costly and unreliable: imagination is helpful for some instances, redundant for many, and misleading for others. This analysis motivates selective invocation and adaptive control.
- Always-on world-model use explores every instance and action branch, incurring substantial cost while producing redundant or noisy observations.
- The analysis categorizes instances as helpful, misleading, or unnecessary according to whether imagined views improve, harm, or do not affect correctness.
- 54% of SAT-Real instances are correctly answered without world-model invocation, while imagination is genuinely helpful in only 14%.
- A 4.6% accuracy gain over baseline requires nearly two orders of magnitude more tokens and about 30× higher inference time under always-on imagination.
- The proposed training loop samples rollouts that either call the world model or skip it, then updates the gating policy using QA-model rewards.
4 Adaptive Visual Imagination Control
AVIC adaptively gates and plans visual imagination, selecting targeted world-model trajectories only when additional evidence is useful. AVIC-R further learns this behavior with reinforcement learning from QA correctness and imagination-cost signals.
- Adaptive Visual Imagination Control: AVIC selectively invokes the world model when additional visual evidence is likely useful, replacing always-on imagination with adaptive test-time control.The framework explicitly gates world-model usage and plans instance-dependent imagination.
- Policy gating with test-time scaling: The policy maps the observation, question, and answer set to either skip or call_wm, together with a short discrete egocentric action plan.Independent policy samples can be aggregated by majority vote to reflect uncertainty about whether imagination is necessary.
- Action execution and trajectory selection: When imagination is requested, the world model renders candidate trajectories, and a verifier evaluates each trajectory as a coherent unit before downstream reasoning.Trajectory-level verification preserves temporal and geometric consistency rather than scoring isolated keyframes.
- Final prediction: The vision-language reasoner predicts from the original observation plus selected imagined views, or from the original observation alone when the gate skips imagination.The selected trajectory therefore directly determines whether additional visual evidence enters final answer prediction.
- Learning when and how much to imagine: AVIC-R trains the gating policy end-to-end with reinforcement learning because optimal imagination trajectories lack ground-truth supervision.Rollouts either query the world model or bypass it, are answered by a frozen QA model, and produce GRPO updates while the environment modules remain frozen.
- Reward design: The composite reward uses QA correctness positively while penalizing imagination length, wrong skips, and unparseable outputs.The design discourages both unnecessary long trajectories and collapse toward always skipping under uncertainty.
- GRPO objective: The policy objective uses group-relative advantages, PPO clipping, and KL regularization to a frozen reference, with only LoRA adapters updated.This preserves the base vision-language model’s general capabilities while shaping gating and planning behavior.
5 Experiments
Experiments show that adaptive imagination improves spatial reasoning and navigation while using substantially fewer world-model calls and tokens than fixed strategies. Analyses identify when imagination helps, how much is sufficient, and how reinforcement learning improves gating.
- Main Results: AVIC matches or surpasses always-on imagination while using ∼10% of the tokens and far fewer world-model calls across backbones on SAT-Real.With GPT-4.1, accuracy rises from 74.0% to 79.3%; with o1, it reaches 85.3% (+10.7% over base).
- Main Results: AVIC-R uses a Qwen2.5VL-7B gating policy to outperform proprietary-policy variants on three of four backbones and improves GPT-4o QA accuracy by 8.0%.Lightweight RL lifts the same policy by 6–8 points across InternVL3-14B, GPT-4o, and GPT-4.1 backbones.
- Main Results: AVIC transfers to R2R navigation, achieving higher OSR, SR, and SPL and lower NE than MapGPT with GPT-4o.The policy selects world-model views at each step and concatenates imagined views with original observations for next-action prediction.
- When Imagination Helps: World-model utility is largest for action-conditioned errors (+57.1%) and smaller for dynamics-understanding errors (+28.5%), showing strong instance dependence.Imagination is often unnecessary when a question only requires symbolic reinterpretation of the current view.
- How Much Imagination: One imagined view raises accuracy from 74.0% to 76%, while a second raises it to 80%; further views provide no gains and can degrade performance.Additional rollouts may introduce rendering artifacts and redundant content that confuse the downstream reasoner.
- Adaptive Scaling: AVIC-R calls the world model on 92% of questions and uses 3.60 views per question, fewer than selective baselines and less than half of MindJourney’s 8.90.Its call rate varies by category, reaching 100% on EgoAct and 78.8% on Perspective.
6 Conclusion
The paper presents adaptive test-time scaling for visual spatial reasoning, selectively controlling when and how much to imagine. Across spatial and navigation benchmarks, AVIC and AVIC-R reduce computation while maintaining competitive or state-of-the-art performance.
- AVIC selectively decides when and how much visual imagination to invoke at inference time.
- AVIC-R trains the gating policy end-to-end with lightweight reinforcement learning using QA-correctness and world-model-cost signals.
- Across spatial reasoning and embodied navigation benchmarks, the framework achieves competitive or state-of-the-art results while reducing world-model calls, tokens, and inference time.
- AVIC-R with a small open-source policy outperforms pipelines using proprietary policy backbones.
- World-model imagination is most beneficial for action-conditioned reasoning and generally requires limited, targeted views.
A Implementation Details
The implementation combines LoRA-based GRPO training, balanced data curation, lenient output parsing, and benchmark-specific evaluation. The framework uses a world model to render imagined trajectories and assesses spatial and navigation performance with task-appropriate metrics.
- AVIC-R post-trains Qwen2.5-VL-7B-Instruct with LoRA using online GRPO and per-question reward normalization.The setup uses K=16 sampled rollouts and approximately 5M trainable parameters.
- The world model renders imagined trajectories from candidate action plans, which are evaluated as coherent temporal and geometric sequences.This contrasts with scoring isolated keyframes.
- Training data uses a curated 30/70 mixture of easy-skip and needs-world-model questions, selected through reference-model prescoring.
- A lenient parser salvages semantically recoverable JSON outputs instead of treating all formatting errors as parse failures.
- SAT and MMSI are evaluated with multiple-choice accuracy, while R2R uses Navigation Error, Oracle Success Rate, Success Rate, and Success weighted by Path Length.
B Extra Experiments
Additional experiments examine robustness, computation, generalization, and gating quality. AVIC reduces costly inference substantially and remains comparatively effective under noisier world-model outputs, but current invocation policies still misalign with some reasoning demands.
- Stage-wise Computation Cost: AVIC reduces total inference time from 177.84 to 29.04 seconds and token usage from 162.6 to 7.6 k while improving accuracy by 2.0 points over always-on imagination.Policy overhead is approximately 14.9 seconds, versus approximately 153.7 seconds saved in world-model time.
- Robustness across runs: AVIC-R achieves mean accuracy 69.33 with a sample standard deviation of 0.77 across three independently sampled runs.
- Sensitivity to Errors from World Models: With Cosmos, overall performance slightly degrades, but AVIC retains a +2.7 improvement over the GPT-o1 baseline.The passage attributes robustness to selective gating and reasoning over multiple imagined perspectives.
- Evaluations on Additional Benchmark: AVIC improves performance on MindCube-Tiny, indicating generalization to an additional fine-grained spatial reasoning benchmark.
- Framework Error Analysis: The gating policy calls the world model most often for Egocentric Movement at 82.6% and Action Consequence at 70.2%, but only 26.4% for Goal tasks.
- Framework Error Analysis: Current invocation remains misaligned with manually identified error sources, producing low recall and precision for cases that truly require world-model imagination.
C Impact Statement
The paper argues that indiscriminate visual imagination can be inefficient or harmful because imagined views may be misleading or redundant. It therefore emphasizes adaptive test-time computation for spatial reasoning.
- Indiscriminate visual imagination can be computationally inefficient and harmful when imagined views are misleading or redundant.
- The findings emphasize selectively invoking visual imagination as a way to support effective spatial reasoning.
D Limitations
The framework is evaluated for adaptive imagination control in visual spatial reasoning and short-horizon embodied navigation, with several scope and design boundaries identified.
- The framework focuses on visual spatial reasoning and short-horizon embodied navigation.
- Extending the framework to longer-horizon decision-making, manipulation, and broader multimodal tasks remains a next direction.
- The framework assumes a separate visual world model and a fixed discrete action space.
E License
The paper states that its code and models will be publicly accessible and identifies licensing information for the resources used.
- The authors will make the code and models publicly accessible.
- The paper provides license information for the datasets, code, and models used.
- QWen2.5VL is listed under the Apache-2.0 license.