Vision-Action_Instruction_Tuning.md

December 14, 2024 · View on GitHub

Vision-Action Instruction Tuning

When using the pre-training model in a specific downstream task, i.e. robot manipulation tasks, we need to further do instruction tuning using collected demonstrations in the new environment setting. In this repo, we provide instructions to reproduce our model for the RLBench Benchmark.

Our code is built mainly on PerAct and RLBench, so make sure to cite them if you find the code useful!

The following content is divided into 4 parts:

Simulation Environment Installation
Preparing Simulation Data
Instruction Tuning
Inference

Simulation Environment Installation

1. You can still use the `llarva` set in pre-training.

conda activate llarva

2. PyRep and Coppelia Simulator

Follow instructions from the official PyRep repo; reproduced here for convenience:

PyRep requires version 4.1 of CoppeliaSim. Download:

Once you have downloaded CoppeliaSim, you can pull PyRep from git:

cd LLARVA/sim
git clone https://github.com/stepjam/PyRep.git
cd PyRep

Add the following to your ~/.bashrc file: (NOTE: the 'EDIT ME' in the first line)

export COPPELIASIM_ROOT=<EDIT ME>/PATH/TO/COPPELIASIM/INSTALL/DIR
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$COPPELIASIM_ROOT
export QT_QPA_PLATFORM_PLUGIN_PATH=$COPPELIASIM_ROOT

Remember to source your bashrc (source ~/.bashrc) or zshrc (source ~/.zshrc) after this.

Warning: CoppeliaSim might cause conflicts with ROS workspaces.

Finally install the python library:

pip install -r requirements.txt
pip install .

You should be good to go! You could try running one of the examples in the examples/ folder.

If you encounter errors, please use the PyRep issue tracker.

3. RLBench

LLARVA uses my RLBench fork.

cd LLARVA/sim
git clone -b LLARVA https://github.com/Dantong88/RLBench # note: 'LLARVA' branch

cd RLBench
pip install -r requirements.txt
python setup.py develop

For running in headless mode, tasks setups, and other issues, please refer to the official repo.

4. YARR

LLARVA uses my YARR fork.

cd LLARVA/sim
git clone -b LLARVA https://github.com/Dantong88/YARR # note: 'LLARVA' branch

cd YARR
pip install -r requirements.txt
python setup.py develop

5. LLARVA-sim

cd LLARVA/sim
pip install pip==24.0 # lower version of pip to support omegaconf==2.0.6
pip install -r requirements.txt

Prepare Simulation Data

Then you need to generate the simulation demonstrations, which are used in our vision-action instruction tuning.

1. Generate RLBench Demos

cd LLARVA/sim/RLBench/tools
export SIM_ROOT=LLARVA/sim
python dataset_generator.py --tasks=sweep_to_dustpan_of_size \
                            --save_path=$SIM_ROOT/data/val \
                            --image_size=128,128 \
                            --renderer=opengl \
                            --episodes_per_task=10 \
                            --processes=1 \
                            --all_variations=True

This example command will generate 10 demos in $SIM_ROOT/data/val, note that you will likely need to generate more demos for training.

2. Adapt the Format (Skip if you only want to test the model)

Before using the demos to do vision-action tuning, you should first transfer the format to get annotations. We provide a script that can be used as follows:

cd /LLARVA/sim
python generate_vision-action-tuning_anns.py --data-path data/val --save-path data/anns/train.json --selected-task sweep_to_dustpan_of_size

You will now see the file train.json in data/anns folder, which can further be used for vision-action instruction tuning.

We also provide one example of the generated demos 800 sweep_to_dustpan_of_size and its processed annotations, for two versions (with/without visual traces). You can download these files and put them in the following structure:

LLARVA/sim
│ 
└── data
│   ├── anns
│   │   ├── sweep_to_dustpan_of_size
│   │       └── (our pre-processed vision-action annotations instruction annotations, including both with/without visual trace)
│   │
│   └── val
│   │   └── (10 demos preciously generated)
│   │
│   └── sweep_to_dustpan_of_size 
│       └── (demos download by link)         
│
└── generate_vision-action-tuning_anns.py
    ...

Instruction Tuning

After getting the training annotations, you can follow these steps to adapt the pre-trained model to a specific downstream task:

1. Put the Pre-training Model

Copy the pre-trained model to the output folder, for example:

cd LLARVA/sim
mkdir output
cd output
mkdir llava-lora-instruction-tuning-sweep_to_dustpan_of_size
cp -r the/path/pretrained_model  llava-lora-instruction-tuning-sweep_to_dustpan_of_size

2. Fix the Package Incompatibilities.

We found some incompatibilities of the deepspeed and transformers package with our fine-tuning code, so you need to go to your conda llarva environment (commonly in anaconda3/envs/llarva/lib/python3.10/sitepackages/), then manually modify the source package as follows:

Set load_module_strict = False, load_optimizer_states=False, load_lr_scheduler_states=False in transformers/integrations/deepspeed.py
Comment out the whole try-except block in transformers/generation/configuration_utils.py
Set self._load_optimizer_and_scheduler(None) at transformers/trainer.py, and add resume_from_checkpoint = None after that.

3. Launch the instruction tuning.

cd LLARVA
bash scripts/v1_5/vision-action_instruction_tuning_rlbench.sh

4. Merge the Lora Weight

After training, you should first merge the lora weights by running:

python scripts/merge_lora_weights.py --model-path 'path/to/your/lora-weights' --model-base 'lmsys/vicuna-7b-v1.5' --save-model-path 'your/path'

Note that your save-model-path should include the word llava, otherwise, you might get an error.

We release our weights for the "meat off grill" task in RLBench as follows (these are the merged final weights, i.e. you do not need to run the above command to merge it before using it in inference):

Lora Weight

Model	Size	Tuning Set	Backbone	Download
LLARVA	7B	800 episodes of sweep_to_dustpan_of_size in RLBench	Vicuna-7B	Model

Inference

To test the model on RLBench, just run

cd LLRVA/sim
export SIM_ROOT=LLRVA/sim
python eval.py \
rlbench.tasks=[sweep_to_dustpan_of_size] \
rlbench.demo_path=$SIM_ROOT/data/val \
framework.eval_from_eps_number=0 \
framework.eval_episodes=10 \
rlbench.episode_length=150 \
framework.gpu=4 \
method.ckpt=path/to/llava-sweep_to_dustpan_of_size_merged # this is the download merged ckpt

Citations

If you find our work inspiring or use our codebase in your research, please consider giving a star ⭐ and a citation.

LLARVA

@misc{niu2024llarva,
      title={LLARVA: Vision-Action Instruction Tuning Enhances Robot Learning}, 
      author={Dantong Niu and Yuvan Sharma and Giscard Biamby and Jerome Quenum and Yutong Bai and Baifeng Shi and Trevor Darrell and Roei Herzig},
      year={2024}
}

PerAct

@inproceedings{shridhar2022peract,
  title     = {Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulation},
  author    = {Shridhar, Mohit and Manuelli, Lucas and Fox, Dieter},
  booktitle = {Proceedings of the 6th Conference on Robot Learning (CoRL)},
  year      = {2022},
}

RLBENCH

@article{james2020rlbench,
  title={Rlbench: The robot learning benchmark \& learning environment},
  author={James, Stephen and Ma, Zicong and Arrojo, David Rovick and Davison, Andrew J},
  journal={IEEE Robotics and Automation Letters},
  volume={5},
  number={2},
  pages={3019--3026},
  year={2020},
  publisher={IEEE}
}