Openvla-oft¶

1. Environment Setup¶

The conda environment for openvla-oft with RoboTwin is identical to the official openvla-oft environment for the ALOHA part. Please follow the (openvla-oft official documentation) to install the environment and directly overwrite the RoboTwin environment.

conda activate RoboTwin
# Install PyTorch
# Use a command specific to your machine: https://pytorch.org/get-started/locally/
pip3 install torch torchvision torchaudio

# Clone openvla-oft repo and pip install to download dependencies
git clone https://github.com/moojink/openvla-oft.git
cd openvla-oft
pip install -e .

# Install Flash Attention 2 for training (https://github.com/Dao-AILab/flash-attention)
#   =>> If you run into difficulty, try `pip cache remove flash_attn` first
pip install packaging ninja
ninja --version; echo $?  # Verify Ninja --> should return exit code "0"
pip install "flash-attn==2.5.5" --no-build-isolation

Note!
If you encounter problems on diffusers, try pip install diffusers==0.33.1

2. Collect RoboTwin Data¶

See RoboTwin Tutorial (Usage Section) for more details.

3. Generate RLDS Data¶

RLDS dataset is the data format required for Openvla-oft training.

use RoboTwin data generation mechanism to generate data.
Then convert the raw data to the aloha format that openvla-oft accepts:

bash preproces_aloha.sh

Then transform the data to tfds form and register the tfds form dataset in your device: e.g.:

python -m datasets.move_can_pot_builder

After converting to RLDS, register the dataset (which, for example, would be called aloha_move_can_pot_builder) with our dataloader by adding an entry for it in configs.py (here), transforms.py (here), and mixtures.py (here).Details in Openvla-oft official documentation

4. Finetune model¶

bash finetune_aloha.sh

By default, the training process will not save merged weights. So you need to run merge_lora.sh to merge lora weights if you want to use the checkpoint. If some .py files miss in the merged checkpoint, just copy them from the original checkpoint.

5. Eval on RoboTwin¶

example usage

bash eval.sh ${task_name} ${task_config} ${checkpoint_path} ${seed} ${gpu_id} ${unnorm_key}
# Example: bash eval.sh move_can_pot demo_randomized ckpt_path 0 5 aloha_move_can_pot_builder

The evaluation results, including videos, will be saved in the eval_result directory under the project root.