TokenPB: Enabling VLA Action Self-Verification via VLM Token Probability Bucketing

Chen Zhao; Zhuoran Wang; Haoyang Li; Shifeng Bao; Guanlin Li; Youhe Feng; Yang Li; Jie Tang

cs.CVMar 18, 2026

TokenPB: Enabling VLA Action Self-Verification via VLM Token Probability Bucketing

Authors: Chen Zhao, Zhuoran Wang, Haoyang Li, Shifeng Bao, Guanlin Li, Youhe Feng, Yang Li, Jie Tang, +1 more

Organizations: School of Information, Renmin University of China · Key Laboratory of Data Engineering and Knowledge Engineering, Beijing, China · Beijing University of Posts and Telecommunications · Tsinghua University

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 $\mathcal{L}_{\text{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.

Figures & tables

Figure 1: Bucket co-training improves optimization and success. Left: adding $\mathcal{L}_{\text{bucket}}$ alongside the flow-matching loss $\mathcal{L}_{\text{fm}}$ accelerates convergence (lower $\mathcal{L}_{\text{fm}}$ at the same step). Right: at every checkpoint, joint training ( $\mathcal{L}_{\text{fm}}+\beta\,\mathcal{L}_{\text{bucket}}$ ) yields higher downstream success rates than the $\mathcal{L}_{\text{fm}}$ -only baseline.

Figure 2: Overview of the TokenPB framework. Left: the VLM backbone learns a mapping from observation states to action token bucket distributions. Right: at test time, the VLA generates multiple candidate action chunks via the diffusion action expert, then scores each candidate by its alignment with the action token bucket probability distribution.

Benchmark	Model	Success Rate (%)
LIBERO		Spatial	Object	Goal	Long	Avg.
	Qwen2.5-VL-GR00T	94.6	97.2	93.8	90.2	93.9
	Qwen2.5-VL-GR00T+TokenPB	96.6	98.8	97.0	94.6	96.8
	InternVL3.5-GR00T	95.8	96.0	94.2	92.8	94.7
	InternVL3.5-GR00T+TokenPB	97.6	99.4	97.2	96.8	97.8
	$\pi_{0.5}$	97.8	98.4	96.6	93.0	96.5

Table 1: Comparison with Single-Sample vs. Test-Time Inference. All models sharing the same backbone are co-trained with both bucket-token and flow-matching objectives using a single checkpoint; configurations differ only at inference time. “Model-GR00T”: single sample (

M{=}1

) without verification. “Model-GR00T+TokenPB”: multi-candidate verification (

M{=}16

). For Real-World, “Seen” denotes in-distribution tasks and “Unseen” denotes held-out instructions and object configurations not seen during training. Detailed training settings are provided in Table 8 .

Figure 3: Inference time comparison with RoboMonkey and MG-Select.

	LIBERO (%)					SimplerEnv-WidowX (SR)
Method	Spatial	Object	Goal	Long	Avg.	Spoon on Towel	Carrot on Plate	Stack Cubes	Eggplant in Basket	Avg.
V-GPS	–	–	–	–	–	0.46	0.16	0.07	0.84	0.38
RoboMonkey	–	–	–	56.5	–	0.68	0.72	0.44	0.18	0.51
MG-Select (single)	97.4	95.4	95.6	79.6	92.0	0.67	0.71	0.42	0.08	0.47
MG-Select	97.2	98.0	94.5	82.7	93.1	0.69	0.75	0.43	0.14	0.50
AsyncVLA	98.0	98.2	96.4	96.0	97.4	0.71	0.67	0.58	0.88	0.71

Table 2: Comparison with baseline verification methods on LIBERO and SimplerEnv-WidowX. Original baseline rows retain results from their official publications; only LIBERO-Long is available for the original RoboMonkey row, and V-GPS is goal-conditioned. Added rows report MG-Select single-sample inference and a reproduced RoboMonkey-Qwen2.5 verifier with the shared

\pi_{0.5}

single-sample baseline. Reproduction details are in Appendix I.1 .

Method	Task 1	Task 2	Task 3	Task 4	Task 5	Avg. Len.
MoDE	96.2	88.9	81.1	71.8	63.5	4.01
MoDE + RoVer	97.1	90.9	82.5	74.9	66.6	4.12 (+0.11)
$\pi_{0.5}$	93.5	86.7	79.6	75.3	71.1	4.07
$\pi_{0.5}$ + TokenPB (Ours)	95.4	90.6	85.3	79.6	75.8	4.27 (+0.20)

Table 3: Comparison with baseline verification methods on CALVIN ABC

\to

D. All baseline results are from their respective official publications.

Figure 4: Test-time cost–performance trade-off on LIBERO-Long as $M$ varies from 1 to 64. Left: Extra inference time (blue, left axis) remains nearly constant across all $M$ . The gray dashed line shows raw GPU memory without optimization; by reusing the observation-prefix KV cache (red), memory grows only to $\sim$ 32 GB even at $M{=}64$ —roughly half the raw cost. Middle: Success rate versus $M$ . TokenPB (blue) improves steadily, whereas random choice (purple) and majority voting (yellow) plateau. Right: Action RMSE versus $M$ . TokenPB drives a monotonic decrease, confirming that selected actions become progressively closer to the ground truth. We adopt $M{=}16$ as the default.

Configuration	LIBERO (%)					RoboTwin2.0 (%)
	Spatial	Object	Goal	Long	Avg.	Blocks Ranking		Open Laptop		Shake Bottle
	Spatial	Object	Goal	Long	Avg.	Easy	Hard	Easy	Hard	Easy	Hard
No Bucket + Single ( $M{=}1$ )	92.4	95.6	91.8	87.6	91.9	32	7	33	13	79	16
No Bucket + Random Select ( $M{=}16$ )	92.0	95.2	91.4	87.8	91.6	30	6	31	12	77	15
Co-trained + Single ( $M{=}1$ )	93.6	96.2	92.8	90.2	93.2	38	11	39	18	81	21
Full TokenPB ( $M{=}16$ )	96.6	98.8	97.0	94.6	96.8	52	24	55	27	86	37

Table 4: Disentangling co-training representation benefit from test-time verification benefit. All variants use Qwen2.5-VL-3B as the VLM backbone. “Co-trained + Single (

M{=}1

)” corresponds to the GR00T baseline in Table 1 . “No Bucket” removes the auxiliary token-bucket cross-entropy

\mathcal{L}_{\text{bucket}}

during training (the action expert is still trained with

\mathcal{L}_{\text{fm}}

); consequently, no learned bucket-likelihood scorer is available at test time and multi-sample selection reduces to random choice. “Full TokenPB” is the complete framework with co-training and sampling-and-selection (

M{=}16

Metric	Adjust Bottle		Blocks Ranking RGB		Open Laptop		Shake Bottle
Metric	TokenPB	Baseline	TokenPB	Baseline	TokenPB	Baseline	TokenPB	Baseline
Avg. Success Steps ( $\downarrow$ )	107.3 (-45.2)	152.5	754.0 (-282.9)	1036.9	142.3 (-17.1)	159.4	95.3 (-54.4)	149.7
Recovery Attempts ( $\uparrow$ )	1.1 (+0.7)	0.4	3.3 (+2.6)	0.7	2.9 (+2.6)	0.3	1.3 (+0.6)	0.7
Env. Collision Before Grasp (%, $\downarrow$ )	3 (-43)	46	14 (-15)	29	0 (-14)	14	4 (-34)	38
Success Rate (%, $\uparrow$ )	26 (+10)	16	24 (+13)	11	27 (+9)	18	37 (+16)	21

Table 5: Fine-grained trajectory metrics on RoboTwin2.0 (Hard). TokenPB-selected actions are compared against single-sample ones from the same co-trained checkpoint. Recovery Attempts : higher values indicate more retry attempts when the initial action fails, reflecting stronger error-correction behavior. Env. Collision Before Grasp : lower values indicate fewer irrelevant collisions with the environment before grasping the target object.

Figure 5: Verification effectiveness analysis. Left: TokenPB maintains near-constant success rate under noise injection compared to random choice. Middle: Bucketing probability scores of successful trajectories are significantly higher than failed ones, and the score magnitude tracks task success rate across suites. Right: Task success rate saturates beyond 1024 buckets.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Qualitative visualization of TokenPB verification across diverse tasks. Each panel shows a task scene with candidate actions and their bucketing probability scores. Candidates with higher scores (shown in blue boxes with check marks) are selected for execution, while lower-scored candidates (red boxes with crosses) are rejected. TokenPB correctly assigns the highest probability to semantically appropriate actions: (a) placing blocks in the correct spatial order in RoboTwin, (b) grasping the target object in LIBERO, (c) selecting the correct objects in a complex real-world instruction task, and (d) pushing the correct blocks in a long-horizon task.

Training Configuration	$\\|\Delta\theta\\|/\\|\theta_{0}\\|$	Cos ( $\nabla_{\text{fm}}$ , $\nabla_{\text{bucket}}$ )	LIBERO SR (%)
Flow-matching Loss only	0.0081	–	91.9
Bucket Loss only	0.0063	–	–
Co-training ( $\beta{=}1$ )	0.0092	+0.31	95.4
Co-training ( $\beta{=}3$ , default)	0.0107	+0.28	96.8
Co-training ( $\beta{=}5$ )	0.0126	+0.24	96.2
Co-training ( $\beta{=}20$ )	0.0168	+0.19	94.7

Appendix

Table 6: Multi-objective training analysis. All configurations use Qwen2.5-VL-3B trained for 64 epochs on LIBERO.

\|\Delta\theta\|/\|\theta_{0}\|

: relative parameter displacement. Cosine(

\nabla_{\text{fm}}

\nabla_{\text{bucket}}

): average cosine similarity between the gradients of the two objectives during training.

Category	Parameter	Value
Model Size	Number of layers	16
	Hidden size	1536
	Number of attention heads	32
	Parameter size	$\sim$ 500M
Action Config	De-noising step	4
Action Config	Number of candidate action chunks $M$	16

Appendix

Table 7: Action Expert Model Parameters.

Category	Parameter	Real-World	LIBERO	CALVIN	RoboTwin2.0 *
Dataset Statistics	Total episodes	2005	1693	17870	2500
	Total frames	366152	273465	1071743	549787
	Task	4	40	389	50
	Instruction type	741	40	389	2420
	FPS	5	10	10	50
Training Parameters	Batch size	144	144	144	120

Appendix

Table 8: Training Configuration on LIBERO, CALVIN, RoboTwin2.0 and Real-world benchmark.

Task	$n$	Baseline ( $\mu\pm\sigma$ )	TokenPB ( $\mu\pm\sigma$ )	TokenPB 95% CI	$p$ -value
Push Blocks	33	$56.2\pm 18.4$	$75.4\pm 15.7$	[70.1, 80.8]	$<0.001$
Clean Table	24	$52.5\pm 21.6$	$72.5\pm 17.3$	[65.4, 79.6]	$<0.005$
Pick & Place	30	$45.5\pm 22.1$	$75.0\pm 16.8$	[68.8, 81.2]	$<0.001$
Hang Cups	24	$43.7\pm 5.8$	$62.5\pm 5.2$	[60.4, 64.6]	$<0.001$
Overall	111	$49.5\pm 18.7$	$71.4\pm 15.1$	[68.6, 74.2]	$<0.001$

Appendix

Table 9: Statistical analysis of Real-World (In-Distribution) results.

n

: total evaluation rollouts per task.

\mu\pm\sigma

: mean success rate (%)

\pm

standard deviation across rollouts. CI: 95% bootstrap confidence interval (1000 resamples).

p

: paired

t

-test

p

-value comparing TokenPB vs. single-sample baseline. All models use Qwen2.5-VL-3B.

Task	$n$	Baseline ( $\mu\pm\sigma$ )	TokenPB ( $\mu\pm\sigma$ )	TokenPB 95% CI	$p$ -value
Push Blocks	54	$5.0\pm 8.7$	$45.0\pm 20.4$	[39.5, 50.5]	$<0.001$
Clean Table	18	$8.0\pm 13.8$	$24.0\pm 18.2$	[15.5, 32.5]	$<0.01$
Pick & Place	54	$28.0\pm 19.5$	$36.0\pm 17.8$	[31.3, 40.7]	$<0.01$
Hang Cups	18	$32.0\pm 5.1$	$40.0\pm 4.5$	[37.9, 42.1]	$<0.001$
Overall	144	$18.3\pm 16.8$	$36.3\pm 18.9$	[33.2, 39.4]	$<0.001$

Appendix

Table 10: Statistical analysis of Real-World (Unseen) results. Each of the 8 unseen instructions is evaluated with 18 independent rollouts (

n=144

total). All models use Qwen2.5-VL-3B.

Figure 7: Partial real machine environment and instruction diagram. Our testing is based on the xArm robotic arm and DaHuan gripper. In real-world we collect 2000 trajectories and select for testing in this work.

Figure 8: Examples of four discrete action token representation methods. Assume VLM tokenizer encode ‘\n’ to 198, ‘0’ to 14, ‘1’ to 15, ‘2’ to 16 and so on.

LIBERO	Spatial	Object	Goal	Long	Avg.
Qwen2.5-VL-Diffusion	94.6	97.2	93.8	90.2	93.9
Action Bins	92.8	95.6	94.4	89.0	93.0 $\downarrow$
FAST	95.0	97.2	93.6	91.8	94.4 $\uparrow$
VLA-0	96.8	95.4	96.0	91.4	94.9 $\uparrow$
Textual FAST	96.6	98.8	97.0	94.6	96.8 $\uparrow$

Appendix

Table 11: Impact of tokenization methods on LIBERO (%). All variants use Qwen2.5-VL-3B with TokenPB (

M{=}16

). The first row (Qwen2.5-VL-Diffusion) reports the single-sample baseline without verification.

\uparrow

\downarrow

indicate improvement/degradation relative to this baseline.

Models	Spatial	Object	Goal	Long	Avg.
Diffusion Policy [ 5 ]	78.3	92.5	68.3	50.5	72.4
HybridVLA [ 21 ]	93.6	90.6	90.2	84.4	90.2
FlowVLA [ 27 ]	93.2	95.0	91.6	72.6	88.1
SmolVLA [ 25 ]	93.0	94.0	91.0	77.0	88.8
OpenVLA-OFT [ 13 ]	97.6	98.4	97.9	94.5	97.1
UniVLA [ 3 ]	95.4	98.8	93.6	94.0	95.5

Appendix

Table 12: Test results on the LIBERO dataset. “Model-FAST”: FAST tokenizer-assisted auto-regression model. “Model-GR00T” and “Model-GR00T+TokenPB” share the same co-trained checkpoint (with both bucket classification and diffusion objectives) and differ only at inference: GR00T uses a single diffusion sample (

M{=}1

), while GR00T+TokenPB applies multi-candidate verification (

M{=}16

). Action chunk size per inference of our model is 16.

Simulation Task	RDT		Pi0		ACT		DP		TokenPB		Qwen-Dif
	Easy	Hard	Easy	Hard	Easy	Hard	Easy	Hard	Easy	Hard	Easy	Hard
Adjust Bottle	81%	75%	90%	56%	97%	23%	97%	0%	97%	26%	96%	16%
Beat Block Hammer	77%	37%	43%	21%	56%	3%	42%	0%	42%	0%	35%	0%
Blocks Ranking RGB	3%	0%	19%	5%	1%	0%	0%	0%	52%	24%	38%	11%
Blocks Ranking Size	0%	0%	7%	1%	0%	0%	1%	0%	22%	0%	16%	0%
Click Alarmclock	61%	12%	63%	11%	32%	4%	61%	5%	52%	22%	50%	19%

Appendix

Table 13: RoboTwin 2.0 Simulation Benchmark.

Method	Adjust Bottle	Move Pill Bottle Cap	Ranking Block RGB
InternVL3.5-2B	14.0	0	3.0
InternVL3.5-2B + TokenPB	22.0	0	5.0
$\pi_{0.5}$	39.0	4.0	11.0
$\pi_{0.5}$ + TokenPB	43.0	11.0	23.0

Appendix

Table 14: RoboTwin2.0 (Hard) results with InternVL3.5-2B and

\pi_{0.5}

backbones. Success rates (%) on selected tasks. TokenPB uses

M{=}16

candidates.

Benchmark	Method	Success rate (%)
LIBERO-Plus	Single sample	76.5
LIBERO-Plus	RoboMonkey	77.6
LIBERO-Plus	TokenPB	76.9
RoboCasa	Single sample	54.9
RoboCasa	RoboMonkey	55.5
RoboCasa	TokenPB	56.4

Appendix

Table 15: Additional baseline results on LIBERO-Plus, RoboCasa, and RoboTwin 2.0. Success rates are reported in percent. The RoboTwin 2.0 result is the mean

\pm

standard deviation over five training seeds.

Method	Spatial	Object	Goal	Long	Avg.
Task success rate
Random selection	92.0	95.2	91.4	87.8	91.6
CQL Q-function	91.8	96.4	90.6	89.2	92.0
TD3+BC Q-function	92.6	96.0	92.4	91.4	93.1
IQL Q-function	93.4	97.8	92.2	91.6	93.7
RoboMonkey-Qwen2.5	97.0	98.4	95.8	92.4	95.9

Appendix

Table 16: Shared-sampler comparison with learned Q-functions on LIBERO (%). All methods select among

M{=}16

candidates. The upper block reports task success; the lower block reports oracle-best candidate selection accuracy on the diagnostic subset. Q-functions use replay-derived rewards, RoboMonkey-Qwen2.5 uses additional preference pairs, and TokenPB uses demonstrations alone.

Component / candidate count	TokenPB	RoboMonkey
Resident sampler	18.2	18.2
Additional verifier	0.0	15.9
Scoring allocation, $M{=}1$	0.1	1.2
Scoring allocation, $M{=}16$	1.1	2.1
Scoring allocation, $M{=}32$	2.7	3.3
Total, $M{=}1$	18.3	35.3

Appendix

Table 17: GPU memory scaling with Qwen2.5-VL-3B (GB). RoboMonkey uses an external Qwen2.5-VL-3B verifier with prefix sharing enabled. The scoring allocation is the total minus the reported resident sampler and, where applicable, external-verifier footprints. Unmeasured entries are left blank.

Backbone / setting	Verif. acc.	Recall acc.	Success rate
LIBERO
Qwen2.5-VL-3B	85.5	33.9	96.8
Qwen3.5-4B	86.0	37.1	97.4
InternVL3.5-2B	87.4	32.2	97.8
Qwen2.5-VL-3B + pretraining	88.2	–	98.6
LIBERO-Plus

Appendix

Table 18: Backbone and pretraining comparisons (%). Verification and recall accuracies use diagnostic subsets. Recall was measured only for the LIBERO backbone comparison. The pretraining rows use Qwen2.5-VL-3B and 10K hours of real-robot data to pretrain bucket-token prediction; other rows use the standard training setting.

Category	Qwen2.5-VL-3B	Qwen3.5-4B	InternVL3.5-2B
Background	92.7	96.7	97.6
Light	93.6	95.6	95.1
Language	88.4	86.7	87.0
Noise	74.6	79.1	80.8
Layout	79.7	83.1	79.0
Robot initialization	54.2	55.2	62.5

Appendix

Table 19: LIBERO-Plus success rates by perturbation category (%). All backbones use TokenPB.

Explore similar work

APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies
Jun 10, 2026Kechun Xu, Zhenjie Zhu, Anzhe Chen +2Visuomotor Policy Learning Vision-Language-Action Models
Zhejiang University · Zhejiang Humanoid Robot Innovation Center
AR-VLA: True Autoregressive Action Expert for Vision-Language-Action Models
Mar 10, 2026Yutong Hu, Jan-Nico Zaech, Nikolay Nikolov +6Vision-Language-Action Models Autoregressive Generation
∗INSAIT, Sofia University “St. Kliment Ohridski” · †KU Leuven, Dept. Mechanical Engineering, Research unit Robotics, Automation and Mechatronics · ‡KU Leuven, Dept. Electrical Engineering, Research unit Processing Speech and Images
Not All Features Are Created Equal: A Mechanistic Study of Vision-Language-Action Models
Mar 19, 2026Bryce Grant, Xijia Zhao, Peng WangVision-Language-Action Models Robotics
Case Western Reserve University

Jun 10, 2026J/K move · Enter open · S save

APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies

Kechun Xu, Zhenjie Zhu, Anzhe Chen, Rong Xiong+1

Zhejiang University · Zhejiang Humanoid Robot Innovation Center

Visuomotor Policy Learning Vision-Language-Action Models Instruction Generalization in VLA Models

Open paper AskPDF arXiv

Vision-Language-Action (VLA) models that couple pretrained Vision-Language Models (VLMs) with continuous action experts have achieved strong manipulation performance, yet generalization to out-of-distribution (OOD) language instructions remains poor. A known challenge is the structural imbalance in VLA data, where language is far less diverse than visual and action content, making policies prone to visual shortcuts. While discrete-action methods mitigate this through vision-language co-training, continuous action experts lack such protection: they start from random initialization and learn entirely from imbalanced data, producing noisy gradients that corrupt the VLM and fail to exploit its language capability. We address this from a Bayesian perspective, factorizing the policy into a language-agnostic Vision-Action (VA) prior and a language-conditioned VLA likelihood, and propose APT, a two-stage training method emphasizing Action expert PreTraining. In Stage 1, the action expert is pretrained as a VA prior on vision-action pairs from a frozen VLM, bypassing the language imbalance. In Stage 2, language tokens are injected through a gated fusion mechanism that integrates VLM features while preserving the learned visuomotor prior. APT applies to mainstream VLA architectures, including the

π

and GR00T-style architectures. Comprehensive experiments validate that APT achieves consistent gains on unseen instructions and compositional tasks. Project Page: https://xukechun.github.io/papers/APT/

Abstract

Explore similar work

APT: Action Expert Pretraining Improves Instruction Generalization of Vision-Language-Action Policies

AR-VLA: True Autoregressive Action Expert for Vision-Language-Action Models

Not All Features Are Created Equal: A Mechanistic Study of Vision-Language-Action Models