TACO: Test-time Anti-exploration via pseudo-COunts

TACO (Test-time Anti-exploration via pseudo-COunts) is a novel test-time scaling framework for VLAs which retains the strong generalization capabilities of pretrained VLAs while effectively constraining outputs to the success modes of specific downstream tasks, performing the Anti-Exploration principle in offline RL. By leveraging a lightweight Coin Flipping Network (CFN), TACO acquires accurate measurement of distributional shift with minimal computational overhead, significantly improving performance on out-of-distribution testcases.

• Principled Anti-Exploration: Mitigates inference-time fragility by constraining generated actions to the "success modes" of the downstream task, effectively handling distribution shifts.
• Universal Compatibility: Seamlessly integrates with Flow-Matching (e.g., $\pi_0$, $\pi_{0.5}$), Diffusion (e.g., RDT), and Autoregressive (e.g., OpenVLA) architectures.
• Gradient-Free Steering: Performs Test-Time Scaling (TTS) via a generate-then-verify pipeline without modifying the heavy VLA backbone parameters.
• Efficient Inference: Implements KV Cache Optimization to reuse visual-language representations, reducing inference latency by ~73% compared to the original manner, during parallel sampling.
• High-Fidelity Verification: Utilizes a lightweight Coin Flipping Network (CFN) trained on internal representations with High-Fidelity Feature Search to accurately estimate action reliability.

1. New Perspective on VLA Instability: We diagnose the inference fragility of generative VLAs as an out-of-support problem and propose TACO, the first framework to address this via the Anti-Exploration principle from Offline RL using Test-Time Scaling.
2. Coupled Pseudo-Count Estimator: We introduce an efficient internal representation mechanism coupled with a High-Fidelity Feature Search strategy. This allows the CFN to accurately verify action chunks for denoising-based policies (Flow/Diffusion) that never see clean actions during training.
3. SOTA Performance: Extensive experiments across extensive simulation tasks (RoboTwin, LIBERO, SimplerEnv) and real-world dual-arm manipulation demonstrate that TACO significantly boosts success rates (e.g., +16% in real-world tasks) over strong baselines like $\pi_0$.

• [2025-12] Releasing TACO code and models. See our huggingface collections.

Create a conda env:

conda create -n taco python=3.10 -y
conda activate taco

Install torch (choose the version that suits your environment):

pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0

Install CFN:

cd cfn/
pip install -e .
cd ..

cd third_party/lerobot
pip install -e .

cd src/transformers
pip install -e .

cd ../../../..

cd ./third_party/Robotwin

Please follow Robotwin doc to install this local Robotwin and its requirements.

cd ./third_party/lerobot
pip install -e ".[libero]"

You can create a new conda environment, refer third_party/openvla to install Libero and our modified OpenVLA.

We provide pre-trained base policies and CFN checkpoints on 🤗 Hugging Face, allowing you to directly evaluate TACO without training.

# Download base policy checkpoints & CFN checkpoints
### CFN state dict is saved at `cfns` sub-directory in the repo
hf download rhodes-team-teleai/pi05_TACO_libero_finetuned --local-dir /path/to/your/dir --max-workers 16

1. preparation

Collect Robotwin task dataset, we provide a pipline as an example:

bash ./scripts/robotwin_data/task_dataset_collection.sh
bash ./scripts/robotwin_data/data_trans/rt2-hdf5_2_hdf5_2_lerobot.sh
bash ./scripts/robotwin_data/data_trans/v21_to_v30.sh
bash ./scripts/robotwin_data/data_trans/make_sure_stats.sh

You will get lerobot dataset v3.0 at repo-id=RoboTwin2/demo_clean/${task}_v30. Please refer to Robotwin doc first if you have any questions.

You can fine-tune your own pi0.5:

bash ./third_party/lerobot/scripts/train_pi05.sh

or just use our trained pi0.5 model on Huggingface.

2. Collect internal representation

Modify and run:

bash ./scripts/collect_inernal_representation/pi05_robotwin2/collect.sh

3. Train CFN

Modify and run:

bash ./scripts/train_cfn/train_cfn_example.sh

4. Eval TACO

Modify and run:

bash ./scripts/eval/eval_robotwin2_torch_pi05_taco.sh

1. Collect internal representation

Run:

bash ./scripts/collect_inernal_representation/pi05_libero/collect.sh

2. Train CFN

Modify and run:

bash ./scripts/train_cfn/train_cfn_example.sh

3. Eval TACO

Modify and run:

bash ./scripts/eval/eval_libero_pi05_taco.sh

1. Collect internal representation

Download openvla/modified_libero_rlds from Huggingface.

Modify and run:

bash ./scripts/collect_inernal_representation/openvla_libero/collect.sh

2. Train CFN

Modify and run:

bash ./scripts/train_cfn/train_cfn_example.sh

3. Eval in Libero

Modify and run:

bash ./scripts/eval/libero_openvla_taco/eval_libero_openvla_taco.sh

Evaluation of success rates (%) on Simpler-WindowX tasks. We compare $\pi_0$ + TACO against the base policy and other state-of-the-art methods.

Our method consistently improves performance across both Flow-Matching ($\pi_{0.5}$) and Autoregressive (OpenVLA) architectures on long-horizon tasks.

Evaluation on the RoboTwin 2.0 benchmark. We report the success rate improvement of TACO over the $\pi_{0.5}$ baseline. The table highlights the top 5 tasks with the most significant gains.

TACO/
├── cfn/                         # Coin Flipping implementation
│   ├── cfn_net.py               # CFN model architecture
│   └── feature_dataset.py       # Dataset for training CFN
├── scripts/
│   ├── collect_internal_representation/  # Scripts to collect features
│   ├── train_cfn/               # CFN training scripts
│   └── eval/                    # Evaluation scripts
├── third_party/
│   ├── lerobot/                 # Pi0.5 (LeRobot) implementation & LeRobot-Libero evaluation
│   ├── openvla/                 # OpenVLA implementation
│   └── Robotwin/                # Robotwin environment
└── README.md

If you find TACO useful for your research, please cite our paper:

@article{yang2025taco,
  title={Steering Vision-Language-Action Models as Anti-Exploration: A Test-Time Scaling Approach},
  author={Siyuan Yang, Yang Zhang, Haoran He, Ling Pan, Xiu Li, Chenjia Bai, Xuelong Li},
  journal={arXiv preprint arXiv:2512.02834},
  year={2025}
}

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.

• LeRobot for the Pi0.5 implementation
• OpenVLA for the OpenVLA base model
• Robotwin for the simulation environment
• Libero for the benchmark tasks

For questions or collaborations, please contact:

• Email: breezeyoung9470@gmail.com
• GitHub Issues: Open an issue