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.
State Key Laboratory of Multimedia Information Processing, School of Computer Science, Peking University · The Chinese University of Hong Kong · Simplexity Robotics