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

TL;DR

UAOR presents a training-free, plug-in approach for Vision-Language-Action models that combats perceptual forgetting by reinjecting observation features into the next FFN layer when layer-wise action entropy indicates high uncertainty. By treating FFNs as key-value memories and using Action Entropy $u_t^{(\ell)}$ with threshold $\gamma$, UAOR restores observation influence during inference, reducing action uncertainty via information-theoretic mechanisms and improving the Information Bottleneck objective. Theoretical guarantees are complemented by extensive experiments across LIBERO, SIMPLER, CALVIN, and real-world robotic tasks, showing consistent performance gains with minimal overhead and without extra observation priors or training. This makes UAOR a practical, generalizable enhancement for diverse VLA architectures and robotics applications, enabling more reliable manipulation with modest computational cost.

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.

Figures (7)

• Figure 1: Layer-wise uncertainty of OpenVLA-OFT across four LIBERO task suites. The dashed red line denotes the chosen uncertainty threshold $\gamma$, while the green segment highlights the last 16 layers.
• Figure 2: Layer-wise cross-attention from action tokens to observation, language, and action tokens in OpenVLA-OFT across four LIBERO task suites.
• Figure 3: Overall framework of UaOR. We compute action entropy at layer $\ell$ to estimate uncertainty. If it exceeds a threshold $\gamma$, we reinject observation features, including visual and proprioceptive features (if available), into the next layer's FFN via a key-value retrieval mechanism, where the input hidden states serve as queries and the observation features act as key-value memory.
• Figure 4: Real-World Setup and Results.
• Figure 5: Impact of uncertainty threshold $\gamma$ and blending factor $\alpha$ across four LIBERO task suites.

Theorems & Definitions (6)

• Theorem 3.1: Observation information gain
• Theorem 3.2: Action uncertainty reduction
• Theorem 3.3: Information Bottleneck optimization
• Theorem 3.4: Benefit of uncertainty-triggered reinjection
• Lemma 1.1: Layerwise MI decay in the vanilla stack
• proof