OpenDriveLab/UniVLA

[!IMPORTANT] 🌟 Stay up to date at opendrivelab.com!

:earth_asia: UniVLA

:page_facing_up: Paper | :rocket: Demo Page (Coming Soon)

:black_nib: Qingwen Bu, Y. Yang, J. Cai, S. Gao, G. Ren, M. Yao, P. Luo, H. Li
:e-mail: Primary Contact: Qingwen Bu ([email protected])

:fire: Highlights

• A recipe towards generalist policy by planning in a unified, embodiment-agnostic action space.
• A novel approach for extracting task-centric latent actions from cross-embodiment videos.
• A VLA that achieves state-of-the-art results on multiple benchmarks with compute-efficient training.

Table of Contents

• :movie_camera: Demo
• :loudspeaker: News
• 🤗 Model Zoo
• :video_game: Getting Started
• :fire: Training Recipe
  • Data Preparation
  • Task-centric Latent Action Learning
  • Pretraining of Generalist Policy
  • Post-training for Deployment & Evaluations
    • Real-world Experiment
    • LIBERO
    • CALVIN
    • Room2Room
    • SimplerEnv
• :rocket: UniVLA's Performance
• :pencil: Citation

:movie_camera: Demo

Real-world robot experiments.

:loudspeaker: News

• [2025/05] The code of UniVLA v1.0 is released. Please check it out!

🤗 Model Zoo

:video_game: Getting Started

1. (Optional) We use conda to manage the environment.

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

2. Install dependencies.

# Install pytorch
# Look up https://pytorch.org/get-started/previous-versions/ with your cuda version for a correct command
# Our experiments are conducted with 'torch 2.2.0 + cuda 12.1'
pip install torch torchvision

# Clone our repo and pip install to download dependencies
git clone [email protected]:OpenDriveLab/UniVLA.git
cd univla
pip install -e .

# Install Flash Attention 2 for training (https://github.com/Dao-AILab/flash-attention)
pip install packaging ninja
ninja --version; echo $?  # Verify Ninja --> should return exit code "0"
pip install "flash-attn==2.5.5" --no-build-isolation

:fire: Training Recipe

:zero: Data Preparation

Please refer to this script for an example of how to download datasets from OXE

[optional] Please follow this instruction if you'd like to convert Ego4D data into RLDS format for training UniVLA.

:one: Task-centric Latent Action Learning

We hightly recommond directly using our pre-trained latent action model ckeckpoints to save your time and compute.

[!NOTE] Our latent action model is trained on a comprehensive data collection, encompassing multiple robotic manipulation and navigation datasets from Open X-Embodiment, along with a curated subset of the Ego4D dataset (detailed data construction procedures are provided in the appendix of our paper).

To adapt the model to additional datasets or custom data sources, users may refer to ./prismatic/vla/datasets/rlds/oxe/mixtures.py to either utilize predefined data mixtures or define new ones. Subsequently, the data_mix parameter in the configuration file should be updated accordingly.

The latent action model is implemented based on VQ-VAE. We train the latent action model on the collection of dataset comprising robot manipulation, navigation and human videos. In stage-1 training, we use an overall batch size of 512 and 100k optimization steps to construct the task-irrelevant latent actions:

torchrun --standalone --nnodes 1 --nproc-per-node 8 main.py fit \
    --config config/lam-stage-1.yaml \
    2>&1 | tee lam-stage-1.log

The following stage-2 then focuses on learning task-centric latent actions on the basis of stage-1 results. Please modify the stage_one_ckpt in latent_action_model/config/lam-stage-2.yaml to your local path of stage-1 checkpoint, then run training with:

torchrun --standalone --nnodes 1 --nproc-per-node 8 main.py fit \
    --config config/lam-stage-2.yaml \
    2>&1 | tee lam-stage-2.log

:two: Pretraining of Generalist Policy

• Latent Action Pseudo-Labeling for Policy Optimization: The trained latent action model is employed to generate pseudo-labels for policy optimization via a next-token prediction objective. Specifically, the indices of inferred latent actions in the VQ-VAE codebook are mapped to dedicated tokens in the LLaMA tokenizer, denoted as {ACT_0, ACT_1, ..., ACT_C}.

• Cost-effective Pre-Training: The full-scale pre-training procedure, incorporating both OpenX and Ego4D datasets, was performed using a 32-GPU A100 cluster over 20,000 optimization steps. This training regimen required approximately 960 A100 GPU-hours, representing just 5% of the computational resources utilized by OpenVLA. Furthermore, experiments conducted on the 'Bridge' and 'Human' subsets demanded only 200 GPU-hours, demonstrating substantially reduced computational requirements compared to previous vision-language-action models.

• To initiate pre-training, please refer to the following scipt or simply run bash ./vla-scripts/train.sh:

[!NOTE] For pretraining UniVLA only on BridgeV2 or Human (Ego4D) data, please modify vla.type to prism-dinosiglip-224px+mx-bridge(human) correspondingly. Detailed setups can be found in ./prismatic/conf/vla.py.

### Experiment on a 32-GPU cluster
GPUS_PER_NODE=8  
NNODES=4
MASTER_PORT=${MASTER_PORT:-28596}
MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
RANK=${RANK:-0}

# Run your training script with torchrun
torchrun --nproc_per_node ${GPUS_PER_NODE} --nnodes ${NNODES} --node_rank ${RANK} --master_addr ${MASTER_ADDR} --master_port ${MASTER_PORT} train.py \
                                 --vla.type prism-dinosiglip-224px+mx-oxe-magic-soup-plus \
                                 --run_root_dir "vla_log" \

Once pretraining is complete, convert the UniVLA weights (default 'Prismatic' format) to HuggingFace AutoClasses with:

python vla-scripts/extern/convert_openvla_weights_to_hf.py \
    --openvla_model_path_or_id /path/to/your/pretrained_ckpt_path \
    --ckpt_name /path/to/your/specific_ckpt_name.pt \
    --output_hf_model_local_path /path/to/your/output_model_path

The converted model is then compatible with HF AutoClasses 'AutoModelForVision2Seq'.

:three: Post-training for Deployment & Evaluations

• With the pretrained generalist policy trained to plan over an embodiment-agnostic action space, we then add embodiment-specific action decoder heads for downstream deployment.
• Our action decoder is extremely lightwight with only around 12M parameters. Using parameter efficient fine-tuning with LoRA rank 32, the total trainable parameter is around 123M.

:mechanical_arm: Real-world Experiment

Our guidelines are based on real-device testing conducted on the AgiLex platform. If you have code deployed on other platforms or in different data formats, we welcome pull requests!

We provide a simple guideline to deploy UniVLA on your customized setups.

1) LIBERO

Please first download the LIBERO datasets that we used in experiments

Start training with torchrun:

1. You should first set the pretrained UniVLA and latent action model path in vla_path and lam_path of the training config.
2. Set your local LIBERO dataset path in data_root_dir.
3. You can choose dataset_name from libero_spatial_no_noops, libero_object_no_noops, libero_goal_no_noops, and libero_10_no_noops

We trained on 'Spatial', 'Object' and 'Goal' for 30k steps and 'Long' for 40k steps. Please first modify the max_steps in training config accordingly for reproduction.

# Start training on LIBERO-10(long) with 8 GPUs
torchrun --standalone --nnodes 1 --nproc-per-node 8 finetune_libero.py \
                                 --dataset_name "libero_10_no_noops" \
                                 --run_root_dir "libero_log" \

Once you finished training and get the action decoder and UniVLA backbone, you can start evaluation with:

# Start evaluation on LIBERO-10
# [Optional] Install LIBERO dependencies
pip install -r experiments/robot/libero/libero_requirements.txt

# By default, we test for 50 rollouts every task, totalling 500 independent trials.
python experiments/robot/libero/run_libero_eval.py \
    --task_suite_name libero_10 \   # Choose from [libero_spatial, libero_object, libero_goal, libero_10] 
    --action_decoder_path /path/to/your/action_decoder_path.pt \
    --pretrained_checkpoint /path/to/your/libero_10_finetuned_univla \
    --save_video False    # Whether to save rollout videos \
    --num_trials_per_task 50 \
    --seed 7

2) CALVIN

Please first follow CALVIN to install relavent dependencies and prepare your dataset

1. You should first set the pretrained UniVLA and latent action model path in vla_path and lam_path of the training config.
2. Set your local CALVIN directory path in calvin_root.
3. Start training with torchrun:

torchrun --standalone --nnodes 1 --nproc-per-node 8 finetune_calvin.py \
                                 --vla_path /path/to/your/univla-7b \
                                 --lam_path /path/to/your/lam-stage-2.ckpt \
                                 --calvin_root /path/to/yout/calvin_root_path \
                                 --max_steps 100000 \
                                 --batch_size 8 \
                                 --grad_accumulation_steps 2 \
                                 --window_size 12 \ 
                                 --run_root_dir "calvin_log" 

Start evaluation on CALVIN:

# Mutli-GPU evaluation is supported
torchrun --standalone --nnodes 1 --nproc-per-node 8 experiments/robot/calvin/run_calvin_eval_ddp.py \
    --calvin_root /path/to/yout/calvin_root_path \
    --action_decoder_path /path/to/your/action_decoder_path.pt \
    --pretrained_checkpoint /path/to/your/calvin_finetuned_univla \
    --seed 7

3) Room2Room

[!NOTE] Please refer to this documentation for detailed guidelines.

4) SimplerEnv

Our SimplerEnv evlauation is based on the official repo.

1. Clone and install SimplerEnv dependencies with

# We used the Maniskill3 version
git clone -b maniskill3 https://github.com/simpler-env/SimplerEnv.git

cd SimplerEnv
pip install --upgrade git+https://github.com/haosulab/ManiSkill.git
pip install -e .

2. Add experiments/robot/simpler-bridge/policies/univla to simpler_env/policies, and replace simpler_env/real2sim_eval_maniskill3.py with experiments/robot/simpler-bridge/real2sim_eval_maniskill3.py.

3. Run evaluation on SimplerEnv-Bridge "Put Spoon on Table Cloth" task:

Please refer to experiments/robot/simpler-bridge/eval_simpler_bridge_4task.sh for the evaluation on all tasks.

ckpt_path="/path/to/your/univla-7b-224-sft-simpler-bridge"
action_decoder_path="/path/to/your/univla-7b-224-sft-simpler-bridge/action_decoder.pt"

CUDA_VISIBLE_DEVICES=0 XLA_PYTHON_CLIENT_PREALLOCATE=false python real2sim_eval_maniskill3.py \
    --model="univla" -e "PutSpoonOnTableClothInScene-v1" -s 0 --num-episodes 24 --num-envs 1 \
    --action_decoder_path ${action_decoder_path} \
    --ckpt_path ${ckpt_path} \

:rocket: UniVLA's Performance

[!NOTE] LIBERO Simulation Benchmark Results.

[!NOTE] LIBERO Results with Limited Data. (Models are trained with 10%, 20%, 50%, and the full dataset)

[!NOTE] SimplerEnv evaluation on WidowX Robot tasks. (Averaged across 3 seeds)

We fix a minor bug about input processing, so the UniVLA's results are higher than the numbers reported in our original paper.

[!NOTE] Real-world Experiments.

:pencil: Citation

If you find our code or models useful in your work, please cite our paper:

@article{bu2025univla,
  title={Univla: Learning to act anywhere with task-centric latent actions},
  author={Bu, Qingwen and Yang, Yanting and Cai, Jisong and Gao, Shenyuan and Ren, Guanghui and Yao, Maoqing and Luo, Ping and Li, Hongyang},
  journal={arXiv preprint arXiv:2505.06111},
  year={2025}
}

Acknowledgements

We thank OpenVLA for their open-sourced work!