Sparse Imagination for Efficient Visual World Model Planning

2025-06-02

AI Summary

To address the high computational cost and poor real-time planning capability of visual world models in resource-constrained robotic settings, this paper proposes a sparse rollout framework. Our method introduces three key innovations: (1) a novel dynamic token pruning mechanism that enables adaptive rollout sparsification based on available compute resources; (2) randomized grouped attention, which balances modeling fidelity and computational efficiency; and (3) a lightweight Transformer-based visual world model trained end-to-end with sparse rollout supervision. Evaluated on multiple visual reinforcement learning benchmarks, our approach maintains task performance while significantly accelerating inference—achieving, for the first time, real-time deployment of world models on edge devices. This work establishes a viable pathway for efficient planning in resource-limited embodied agents.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Computer Vision: Diffusion Models for Vision
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Responsible Web: Human-perceived consequences of algorithmic deployment on the web

Abstract

World model based planning has significantly improved decision-making in complex environments by enabling agents to simulate future states and make informed choices. However, ensuring the prediction accuracy of world models often demands substantial computational resources, posing a major challenge for real-time applications. This computational burden is particularly restrictive in robotics, where resources are severely constrained. To address this limitation, we propose a Sparse Imagination for Efficient Visual World Model Planning, which enhances computational efficiency by reducing the number of tokens processed during forward prediction. Our method leverages a sparsely trained vision-based world model based on transformers with randomized grouped attention strategy, allowing the model to adaptively adjust the number of tokens processed based on the computational resource. By enabling sparse imagination (rollout), our approach significantly accelerates planning while maintaining high control fidelity. Experimental results demonstrate that sparse imagination preserves task performance while dramatically improving inference efficiency, paving the way for the deployment of world models in real-time decision-making scenarios.

Problem

• Reducing computational resources for world model planning
• Improving real-time decision-making in resource-constrained robotics
• Maintaining accuracy while accelerating visual world model planning

Innovation

• Sparse Imagination reduces tokens for efficiency
• Randomized grouped attention in transformer model
• Adaptive token processing based on resources