MosaicThinker: On-Device Visual Spatial Reasoning for Embodied AI via Iterative Construction of Space Representation

When embodied AI is expanding from traditional object detection and recognition to more advanced tasks of robot manipulation and actuation planning, visual spatial reasoning from the video inputs is necessary to perceive the spatial relationships of objects and guide device actions. However, existing visual language models (VLMs) have very weak capabilities in spatial reasoning due to the lack of knowledge about 3D spatial information, especially when the reasoning task involve complex spatial relations across multiple video frames. In this paper, we present a new inference-time computing technique for on-device embodied AI, namely \emph{MosaicThinker}, which enhances the on-device small VLM’s spatial reasoning capabilities on difficult cross-frame reasoning tasks. Our basic idea is to integrate fragmented spatial information from multiple frames into a unified space representation of global semantic map, and further guide the VLM’s spatial reasoning over the semantic map via a visual prompt. Experiment results show that our technique can greatly enhance the accuracy of cross-frame spatial reasoning on resource-constrained embodied AI devices, over reasoning tasks with diverse types and complexities.

💡 Research Summary

The paper addresses a critical gap in embodied AI: small on‑device visual language models (VLMs) excel at appearance‑based tasks but struggle with spatial reasoning that requires understanding 3D relationships across multiple video frames. Retraining or redesigning large VLMs for this purpose is impractical for resource‑constrained devices. To bridge this gap without modifying the VLM itself, the authors propose MosaicThinker, an inference‑time framework that augments a VLM with a compact, sparse global semantic map built from fragmented spatial cues extracted from the video stream.

MosaicThinker operates in four stages. First, each video frame is processed by lightweight auxiliary models (object detector, depth estimator, and a 3‑D point‑cloud generator) to obtain per‑frame object poses, sizes, and camera extrinsics. Second, a cross‑frame alignment module matches objects across consecutive frames using a combination of visual feature similarity and 3‑D transformation consistency, thereby registering all per‑frame data into a common global coordinate system. Third, instead of using every frame, MosaicThinker iteratively selects a subset of “key frames” that are most informative for the current query. The selection loop starts with random sampling, computes a query‑relevance score for each frame (based on new object appearances, viewpoint changes, and textual keyword overlap), and refines the sampling distribution over several iterations. This reduces computational load and filters out noisy or irrelevant observations.

From the aligned key frames, a sparse global semantic map is constructed. Unlike dense bird’s‑eye‑view (BEV) maps, the semantic map consists of a grid of cells, each storing a structured token that encodes object ID, 3‑D position, orientation, and size. The sparsity makes the map easily digestible by small VLMs that lack the capacity to process high‑resolution dense images.

The final stage injects the semantic map into the VLM via a visual prompt. The prompt combines a natural‑language query with a rendered view of the semantic map (e.g., “Refer to the map below and tell me which object is to the right of the sneaker from the current camera pose”). The VLM processes the combined input, leveraging its language understanding while grounding its reasoning on the explicit spatial representation supplied by the map. No parameters of the VLM are altered; only the input modality changes.

The authors evaluate MosaicThinker on three on‑device platforms—NVIDIA Jetson Orion, Meta AR Glass, and OnePlus 12R—across eight indoor environments (homes, offices, libraries) and five categories of spatial queries (object‑relation, location identification, camera motion estimation, size comparison, multi‑object chain reasoning). Baselines include standard video‑text models, BEV‑based methods, and approaches that inject 3‑D tokens or depth maps into VLMs. MosaicThinker consistently outperforms these baselines, improving average accuracy from ~68 % to ~92 % (up to a 40 percentage‑point gain) even when the underlying VLM is a modest 7‑billion‑parameter model. The additional computational overhead remains under 12 % of the total inference budget, preserving real‑time performance (≈30 FPS).

Limitations are acknowledged: the quality of auxiliary modules (depth estimation, point‑cloud generation) heavily influences overall performance; dynamic scenes with moving objects or abrupt lighting changes are not fully addressed; and the key‑frame selection policy may need further adaptation for diverse query types. Future work is suggested in three directions: (1) fusing additional modalities such as IMU or LiDAR to improve pose estimation robustness, (2) developing continuous map‑update mechanisms for dynamic environments, and (3) learning query‑conditioned frame‑selection strategies to further reduce latency and improve relevance.

In summary, MosaicThinker demonstrates that sophisticated cross‑frame spatial reasoning can be achieved on low‑power devices by augmenting existing small VLMs with a lightweight, sparsely encoded semantic map constructed at inference time. This training‑free approach opens the door for advanced embodied AI capabilities—such as precise robotic manipulation, AR assistance, and autonomous drone navigation—on platforms that were previously limited to simple appearance‑based perception.