UAOR: Uncertainty-aware Observation Reinjection for Vision-Language-Action Models
Authors: Jiabing Yang, Yixiang Chen, Yuan Xu, Peiyan Li, Zichen Wen, Bowen Fang, Tao Yu, Xiangnan Wu, +9 more
Organizations: School of Artificial Intelligence, University of Chinese Academy of Sciences · New Laboratory of Pattern Recognition (NLPR), Institute of Automation, Chinese Academy of Sciences · Shanghai Jiao Tong University · FiveAges
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 directly augments the hidden states with observation evidence at high-uncertainty layers, enabling more accurate and reliable 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 & tables
Figure 1: Layer-wise uncertainty of OpenVLA-OFT across four LIBERO task suites. The dashed red line denotes the chosen uncertainty threshold γ , 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 ℓ to estimate uncertainty. If it exceeds a threshold γ , 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.
Method
Spatial
Object
Goal
Long
Average
OpenVLA [ 26 ] ( CoRL’24 )
84.7
88.4
79.2
53.7
76.5
TraceVLA [ 52 ] ( ICLR’25 )
84.6
85.2
75.1
54.1
74.8
SpatialVLA [ 40 ] ( RSS’25 )
88.2
89.9
78.6
55.5
78.1
UniVLA [ 10 ] ( RSS’25 )
96.5
96.8
95.6
92.0
95.2
CogVLA [ 31 ] ( NeurIPS’25 )
98.6
98.8
96.6
95.4
97.4
3D-CAVLA [ 5 ] ( arXiv’25 )
98.2
99.8
98.2
96.1
98.1
Table 1: Performance comparison on the LIBERO benchmark. “†” indicates our reproduced results.
Method
Pick
Move
Open/Close
Open
Average
Coke Can
Near
Drawer
and Place
RT-1 [ 8 ] ( arXiv’23 )
85.7
44.2
73.0
6.5
52.4
RT-2-X [ 44 ] ( CoRL’23 )
78.7
77.9
25.0
3.7
46.3
OpenVLA [ 26 ] ( CoRL’24 )
18.0
56.3
63.0
0.0
34.3
CogACT† [ 30 ] ( arXiv’24 )
92.3 ± 0.3
83.7 ± 0.6
72.7 ± 0.2
43.5 ± 1.0
73.1 ± 0.7
w/ UaOR ( Ours )
95.0 ± 0.3
87.1 ± 0.3
73.6 ± 0.4
47.2 ± 0.4
75.7 ± 0.5
Table 2: Performance comparison on the SIMPLER benchmark. “†” indicates our reproduced results.
Method
Success Rate (%)
Avg. Len
1/5
2/5
3/5
4/5
5/5
RoboFlamingo [ 32 ] ( ICLR’24 )
82.4
61.9
46.6
33.1
23.5
2.47
Vidman [ 47 ] ( NIPS’24 )
91.5
76.4
68.2
59.2
46.7
3.42
OpenVLA [ 26 ] ( CoRL’24 )
91.3
77.8
62.0
52.1
43.5
3.27
VLAS [ 51 ] ( ICLR’25 )
87.2
64.2
40.9
28.1
19.6
2.40
LLaVA-VLA† [ 41 ] ( ICRA’26 )
94.4 ± 0.2
82.0 ± 0.8
70.8 ± 0.3
59.4 ± 0.6
48.2 ± 0.4
3.55 ± 0.05
Table 3: Performance comparison on the CALVIN benchmark. “†” indicates our reproduced results.
Figure 4: Real-World Setup and Results.
Extractor
Trigger
Spatial
Object
Goal
Long
Avg.
Baseline
98.2
98.2
97.6
94.2
97.1
Mean Pool
All
98.0
96.8
95.8
94.4
96.3
Random
98.4
97.8
97.8
94.8
97.2
Entropy
98.0
97.8
97.6
93.8
96.8
Attn Retr.
All
97.8
97.6
96.2
95.2
96.7
Random
97.8
97.6
96.4
93.6
96.4
Table 4: Factorial ablation on LIBERO (OpenVLA-OFT): feature extractor (mean pool vs. attention retrieval) × trigger (all / random / entropy).
#
V.
P.
I.
Spatial
Object
Goal
Long
Avg.
1
✗
✗
✗
98.2
98.2
97.6
94.2
97.1
2
✓
✗
✗
98.4
98.0
97.2
94.6
97.1
3
✗
✓
✗
97.4
97.4
97.4
93.4
96.4
4
✗
✗
✓
98.4
98.4
97.0
93.8
96.9
5
✓
✓
✗
99.0
98.4
98.2
96.2
98.0
6
✓
✗
✓
97.6
97.8
96.6
93.4
96.4
Table 5: Reinjection information ablation on LIBERO based on OpenVLA-OFT. V., P., I. denote Vision, Proprio, and Instruction.
Figure 5: Impact of uncertainty threshold γ and blending factor α across four LIBERO task suites.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Benchmark
Base Model
Task / Suite
γ
α
LIBERO
OpenVLA-OFT
Spatial
0.75
0.05
Object
0.80
0.05
Goal
0.75
0.05
Long
0.80
0.05
π0
Spatial
0.20
0.05
Object
0.20
0.05
Appendix
Table 6: UaOR hyperparameters on simulation and real-world benchmarks
Hyperparameter
Value
# GPUs
8 x NVIDIA 4090 (24GB VRAM)
learning rate (LR)
5e-4
total batch size
8 (1 per GPU)
# train steps
150K
input images
1 third-person camera image
input image size
224 x 224 px
Appendix
Table 7: OpenVLA-OFT hyperparameters for real-world fine-tuning.
Hyperparameter
Value
# GPUs
8 x NVIDIA A100 (80GB VRAM)
learning rate (LR)
2e-5
total batch size
8 (1 per GPU)
input images
1 third-person camera image
input image size
224 x 224 px
VLM backbone
Prism-DinoSigLIP-224px
Appendix
Table 8: CogACT hyperparameters for real-world fine-tuning.
Injection Timing
Injection Module
Success Rate (%)
Latency
Overhead
Spatial
Object
Goal
Long
Avg.
-
Baseline (No Injection)
98.2
98.2
97.6
94.2
97.1
0.161s
-
Current Layer ( ℓ )
Self-Attention (SA)
98.2
98.0
97.8
95.8
97.5
0.195s
+21.1%
Current Layer ( ℓ )
Feed-Forward (FFN)
98.6
98.2
98.0
95.8
97.7
0.182s
+13.0%
Next Layer ( ℓ+1 )
Self-Attention (SA)
98.4
98.0
97.8
94.8
97.3
0.170s
+5.6%
Next Layer ( ℓ+1 )
Feed-Forward ( UaOR )
99.0
98.4
98.2
96.2
98.0
0.169s
+5.0%
Appendix
Table 9: Ablation on Injection Timing and Location on LIBERO based on OpenVLA-OFT.
Method
Throughput ↑
Latency ↓
OpenVLA-OFT
49.7 Hz
0.161 s
w/ UaOR
47.3 Hz
0.169 s
Δ
−4.8%
+5.0%
Appendix
Table 10: Inference overhead of UaOR on OpenVLA-OFT (LIBERO-Long, 500 rollouts). Throughput : generated actions per second; Latency : time per step.
Figure 6: Manipulation Visualizations in the LIBERO Simulation Environment. We present the execution processes of OpenVLA-OFT with UaOR across LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-Long, demonstrating its strong performance under diverse instructions and a wide range of tasks. Each row shows a temporally ordered sequence from left to right.
Figure 7: Manipulation Visualizations in the Real-World Environment. We present the execution processes of OpenVLA-OFT with UaOR across four real-world tasks, demonstrating its strong effectiveness and practicality in real-world scenarios. Each pair of rows shows a temporally ordered sequence from left to right.
Vision-Language-Action (VLA) models have shown strong performance in robotic manipulation, but reliable uncertainty quantification remains challenging, particularly under distribution shift. Unlike autoregressive policies, many modern VLA models generate continuous actions through regression or flow-based generation, where explicit predictive probabilities are unavailable. Moreover, stochastic action sampling primarily captures action-generation variability under a fixed model, while failure detection under distribution shift can benefit from capturing uncertainty in the model itself. Motivated by Bayesian perspectives on local model variations, we propose perturbation-based failure detection (PFD), a training-free framework for estimating epistemic uncertainty in VLA models through low-rank weight perturbations. Specifically, we inject random low-rank perturbations into selected transformer weight matrices and estimate epistemic uncertainty from disagreement across perturbed action predictions. Experiments on LIBERO-PRO show that PFD achieves the highest average AUROC and balanced accuracy among the evaluated methods while consistently outperforming stochastic action sampling across distribution shifts. Real-world robot experiments further demonstrate that PFD provides a competitive failure-detection signal under an unseen object shift.
Yousung Lee, Dongsoo Har
Korea Advanced Institute of Science and Technology (KAIST), Daejeon, Republic of Korea.
Vision-Language-Action (VLA) models have recently emerged as a promising paradigm for robotic manipulation, in which reliable action prediction critically depends on accurately interpreting and integrating visual observations conditioned on language instructions. Although recent works have sought to enhance the visual capabilities of VLA models, most approaches treat the LLM backbone as a black box, providing limited insight into how visual information is grounded into action generation. Therefore, we perform a systematic analysis of multiple VLA models across different action-generation paradigms and observe that sensitivity to visual tokens progressively decreases in deeper layers during action generation. Motivated by this observation, we propose DeepVision-VLA, built on a Vision-Language Mixture-of-Transformers (VL-MoT) framework. This framework enables shared attention between the vision foundation model and the VLA backbone, injecting multi-level visual features from the vision expert into deeper layers of the VLA backbone to enhance visual representations for precise and complex manipulation. In addition, we introduce Action-Guided Visual Pruning (AGVP), which leverages shallow-layer attention to prune irrelevant visual tokens while preserving task-relevant ones, reinforcing critical visual cues for manipulation with minimal computational overhead. DeepVision-VLA outperforms prior state-of-the-art methods by 9.0% and 7.5% on simulated and real-world tasks, respectively, providing new insights for the design of visually enhanced VLA models.
Yulin Luo, Hao Chen, Zhuangzhe Wu +10
State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · The Chinese University of Hong Kong · Simplexity Robotics
Vision-language-action (VLA) models have emerged as generalist robotic controllers capable of mapping visual observations and natural language instructions to continuous action sequences. However, VLAs provide no calibrated measure of confidence in their action predictions, thus limiting their reliability in real-world settings where uncertainty and failures must be anticipated. To address this problem we introduce ReconVLA, a reliable conformal model that produces uncertainty-guided and failure-aware control signals. Concretely, our approach applies conformal prediction directly to the action token outputs of pretrained VLA policies, yielding calibrated uncertainty estimates that correlate with execution quality and task success. Furthermore, we extend conformal prediction to the robot state space to detect outliers or unsafe states before failures occur, providing a simple yet effective failure detection mechanism that complements the action-level uncertainty. We evaluate ReconVLA in both simulation and real robot experiments across diverse manipulation tasks. Our results show that conformalized action predictions consistently improve failure anticipation, reduce catastrophic errors, and provide a calibrated measure of confidence without retraining or modifying the underlying VLA.
Lingling Chen, Zongyao Lyu, William J. Beksi
Department of Computer Science and Engineering, The University of Texas at Arlington, Arlington, TX, USA