UAOR: Uncertainty-aware Observation Reinjection for Vision-Language-Action Models

2026-02-20

AI Summary

This work proposes a training-free, plug-and-play mechanism to enhance vision-language-action (VLA) models without requiring additional data, auxiliary modules, or fine-tuning. The approach uniquely leverages the feed-forward network (FFN) layers of large language models as memory carriers and dynamically assesses uncertainty via action entropy. When high uncertainty is detected, critical visual observations are selectively reinjected into subsequent FFN layers through an attention mechanism, thereby strengthening the model’s focus on relevant perceptual inputs. This method consistently improves performance across multiple VLA architectures in both simulated and real-world robotic tasks, achieving notable gains with minimal computational overhead.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Large Multimodal Models (LMMs)
• Humans and AI: Human-Aware Planning and Behavior Prediction

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

Abstract

Vision-Language-Action (VLA) models leverage pretrained Vision-Language Models (VLMs) as backbones to map images and instructions to actions, demonstrating remarkable potential for generalizable robotic manipulation. To enhance performance, existing methods often incorporate extra observation cues (e.g., depth maps, point clouds) or auxiliary modules (e.g., object detectors, encoders) to enable more precise and reliable task execution, yet these typically require costly data collection and additional training. Inspired by the finding that Feed-Forward Network (FFN) in language models can act as"key-value memory", we propose Uncertainty-aware Observation Reinjection (UAOR), an effective, training-free and plug-and-play module for VLA models. Specifically, when the current language model layer exhibits high uncertainty, measured by Action Entropy, it reinjects key observation information into the next layer's Feed-Forward Network (FFN) through attention retrieval. This mechanism helps VLAs better attend to observations during inference, enabling more confident and faithful action generation. Comprehensive experiments show that our method consistently improves diverse VLA models across simulation and real-world tasks with minimal overhead. Notably, UAOR eliminates the need for additional observation cues or modules, making it a versatile and practical plug-in for existing VLA pipelines. The project page is at https://uaor.jiabingyang.cn.

Problem

• Vision-Language-Action models
• observation cues
• auxiliary modules
• data collection
• training overhead

Innovation

• Uncertainty-aware
• Observation Reinjection
• Vision-Language-Action Models
• Action Entropy
• Training-free