TokenPB: Enabling VLA Action Self-Verification via VLM Token Probability Bucketing
cs.CV, cs.RO
Submitted: 2026-03-18
Updated: 2026-10-12
Code: https://github.com/RUCKBReasoning/TokenPB
License: http://creativecommons.org/licenses/by/4.0/
The gist: Test-time sampling can improve Vision-Language-Action (VLA) policies, but only if the model can efficiently select a good action from multiple candidates.
Terminology
Abstract
Test-time sampling can improve Vision-Language-Action (VLA) policies, but only if the model can efficiently select a good action from multiple candidates. Prior work often relies on separately trained external verifiers, adding extra models and inference overhead. We propose Token Probability Bucketing (TokenPB), a self-verification approach that discretizes continuous action chunks into bucket tokens and co-trains a single VLA checkpoint with a joint objective: a flow-matching loss for action generation and an auxiliary bucket loss L bucket that teaches the VLM backbone to model observation-conditioned distributions over bucket-token sequences. At inference, TokenPB samples M candidates, scores each via teacher-forced evaluation under the learned distribution, and selects the best---enabling efficient batched scoring with observation-prefix KV-cache reuse. We provide theory connecting the training objective to a monotonic ranking property and bounds on expected selection regret that separate candidate-set effects from calibration. Across multiple backbones and benchmarks, TokenPB yields consistent success-rate improvements over single-sample inference from the same checkpoint (+2.7 in simulation, +17.1 in real-world, +16.4 under out-of-distribution shifts), and further improves over external-verifier baselines under the same sampling budget. The code is publicly available at https://github.com/RUCKBReasoning/TokenPB.
Sources
- Qwen2.5-VL Technical Report
- GR00T N1: An Open Foundation Model for Generalist Humanoid Robots
- $\pi_0$: A Vision-Language-Action Flow Model for General Robot Control
- Real-Time Execution of Action Chunking Flow Policies
- RT-1: Robotics Transformer for Real-World Control at Scale
- UniVLA: Learning to Act Anywhere with Task-centric Latent Actions
- RoboTwin 2.0: A Scalable Data Generator and Benchmark with Strong Domain Randomization for Robust Bimanual Robotic Manipulation
- EveryDayVLA: A Vision-Language-Action Model for Affordable Robotic Manipulation
- Knowledge Insulating Vision-Language-Action Models: Train Fast, Run Fast, Generalize Better
- Theoretical Benefit and Limitation of Diffusion Language Model
- VLA-0: Building State-of-the-Art VLAs with Zero Modification
- Galaxea Open-World Dataset and G0 Dual-System VLA Model
- Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success
- MolmoAct: Action Reasoning Models that can Reason in Space
- Survey of Vision-Language-Action Models for Embodied Manipulation
- Flow Matching for Generative Modeling
- HybridVLA: Collaborative Diffusion and Autoregression in a Unified Vision-Language-Action Model
- RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation
- FAST: Efficient Action Tokenization for Vision-Language-Action Models
- Large VLM-based Vision-Language-Action Models for Robotic Manipulation: A Survey
Related papers
- Loss Knows Best: Detecting Annotation Errors in Videos via Loss Trajectories
- AnchorWeave: World-Consistent Video Generation with Retrieved Local Spatial Memories
- Benchmarking the Robustness of Foundation Models for Mammography under Domain Shift
- MambaX-Net: Dual-Input Mamba-Enhanced Cross-Attention Network for Longitudinal MRI Segmentation
- TeleOCR: Navigating Document Parsing Across Digital and Camera-Captured Documents
- A Survey on Efficient Vision-Language-Action Models