TTRL-VL: Test-Time Reinforcement Learning for vision language model

✨Getting Started

conda create -y -n ttrl python=3.10
conda activate ttrl
git clone https://github.com/shyammarjit/TTRL_VL.git
cd TTRL_VL/verl
# Install torch (must have 12.4 cuda version) (for KIAC cluster)
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
# OR Install torch (if you have 12.1 cuda version) (for DGXH100 cluster)
pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu121
# Install flash-attention (Try-1)
pip install -U flash-attn
# Install flash-attention (Try-2)
FLASH_ATTENTION_SKIP_CUDA_BUILD=TRUE pip install flash-attn --no-build-isolation
pip install python-dotenv
# Install transformers for Intern-VL model
pip install transformers==4.52.4 
# Install transformers for Qwen model
pip install transformers==4.51.3
pip3 install vllm==0.8.3
pip install qwen_vl_utils
python -m pip install cffi
pip install latex2sympy2 word2number timm librosa soundfile
# Install TTRL framework 
pip install -e.
# Install other packages
pip install -r requirements.txt

Download datasets

1. ImageNet-1K        -> https://huggingface.co/datasets/ILSVRC/imagenet-1k/tree/main)
2. ImageNet-R         ->
3. ImageNet-Sketch    -> 
4. ImageNet-A         ->

Prepare dataset

# STEP-1
# After downloading the dataset please run to generate the training samples 
# (for e.g. 20, 500, or full-scale) and full-scale testing samples
# RUN inside `TTRL_VL` folder
python gen_training_testing.py

# STEP-2
# Thereafter copy the train and test json files inside `TTRL_VL/verl/data` folder
# Keep the same dataset name and put the files inside of that.

# STEP-3
# Thereafter, generate parquet file from json file inside `TTRL_VL/verl/data` folder. 
# RUN inside `TTRL_VL/verl/data` folder
python preprocess.py

RUN TTRL

# default script
bash examples/ttrl/aime.sh


# Note:
# for internVL keep:    actor_rollout_ref.rollout.temperature=1.0 & actor_rollout_ref.rollout.val_kwargs.temperature=0
# for Qwen keep:        actor_rollout_ref.rollout.temperature=1.5 & actor_rollout_ref.rollout.val_kwargs.temperature=0
# Imagenet-a
bash examples/ttrl/imagenet_a_internvl3_20e.sh