You can deploy this model and start earning money today!

merlo00/
OpenVLA-IsaacLab-finetuning

Views

No views yet

Google Colab Finetuning Pipeline for OpenVLA

This repository provides a Google Colab pipeline for fine-tuning OpenVLA (Open Vision-Language-Action) models. Below is a step-by-step guide to understanding and using the code provided in the Colab notebook.

Introduction
OpenVLA folder editing
Package installation
Dataset loading
The finetuning script
Saving checkpoints

Introduction

OpenVLA is a powerful model for vision-language-action tasks. This Colab pipeline allows you to fine-tune OpenVLA on your custom dataset. The pipeline includes OpenVLA folder editing, package installation, dataset loading, the finetuning script and saving the checkpoints.

OpenVLA folder editing

The OpenVLA folder needs to be adapted in the following points. In the openvla/vla/datasets/rlds/oxe/configs.py define your processed images (according to the dataset in use) and the state and action encoding:

python

    "overall900": {                                                                                                
        "image_obs_keys": {"primary": "image_front", "secondary":"image_side", "wrist": None},  
        "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None},
        "state_obs_keys": ["EEF_state", None, "gripper_state"],
        "state_encoding": StateEncoding.POS_EULER,    
        "action_encoding": ActionEncoding.EEF_POS,   
    }, 

In the openvla/vla/datasets/rlds/oxe/transforms.py Define the transformation function to relabel data from your specific dataset into the format required by the OpenVLA model:

python

def overall900_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]:


    trajectory["observation"]["EEF_state"] = trajectory["observation"]["state"][:, :6]
    trajectory["observation"]["gripper_state"] = trajectory["observation"]["state"][:, -1]
    trajectory["language_instruction"] = trajectory["natural_language_instruction"]

    return trajectory

python

 "overall900": overall900_dataset_transform, 

In the openvla/vla/datasets/rlds/oxe/mixtures.py Add your dataset to the overall mixture (by default, the weight is set to 1.0):

python

    "overall900": [
        ("overall900", 1.0),
    ],

Package installation

First, ensure your edited and zipped openvla folder is loaded correctly. Run the following code in your Colab environment:

python

# Extract and unzip the OpenVLA folder
zip_path = "/content/openvla.zip"
!unzip -o "$zip_path" -d /content

# Install PyTorch with CUDA (update the CUDA version if necessary, based on Colab support)
!pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118

# Minimum requirements: timm==0.9.10, tokenizers==0.19.1, torch>=2.2.0, torchvision>=0.16.0, transformers==4.40.1

%cd openvla
!pip install -e .

# Install Flash Attention 2 and Ninja
!pip install packaging ninja==1.11.1.2
!ninja --version  # Verify Ninja installation

!pip install "flash-attn==2.5.5" --no-build-isolation
!pip install bitsandbytes==0.45.0
!pip install accelerate==0.24.1

Dataset loading

Datasets need to be converted in the Reinforcement Learning Dataset Specification (RLDS) format prior to be processed by the OpenVLA model. This conversion is possible with the help of the RLDS dataset converter made at disposal by one of the Google DeepMind research engineers (available at the following link: https://github.com/kpertsch/rlds_dataset_builder). After cloning the folder, the converting script can be found at the path rlds_dataset_builder/overall_900/overall_900_dataset_builder.py . Here is the script for convereting the dataset starting from episodes saved in the Numpy format (also find a plausible but not fixed structure for these datasets):

python

from typing import Iterator, Tuple, Any
import glob
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_hub as hub

class Overall900(tfds.core.GeneratorBasedBuilder):
    """DatasetBuilder for example dataset."""

    VERSION = tfds.core.Version('1.0.0')
    RELEASE_NOTES = {
        '1.0.0': 'Initial release.',
    }

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder-large/5")

    def _info(self) -> tfds.core.DatasetInfo:
        """Dataset metadata (homepage, citation,...)."""
        return self.dataset_info_from_configs(
            features=tfds.features.FeaturesDict({
                'steps': tfds.features.Dataset({
                    'observation': tfds.features.FeaturesDict({
                        'image_front': tfds.features.Image(
                            shape=(224, 224, 3),
                            dtype=np.uint8,
                            encoding_format='png',
                            doc='Main camera RGB observation.',
                        ),
                        'image_side': tfds.features.Image(
                            shape=(224, 224, 3),
                            dtype=np.uint8,
                            encoding_format='png',
                            doc='Side camera RGB observation.',
                        ),
                        'image_hand': tfds.features.Image(
                            shape=(224, 224, 3),
                            dtype=np.uint8,
                            encoding_format='png',
                            doc='Hand camera RGB observation.',
                        ),
                        'state': tfds.features.Tensor(
                            shape=(7,),
                            dtype=np.float32,
                            doc='Robot state, consists of [3x robot displacements, 3x rotations, and gripper closedness].',
                        ),
                    }),
                    'action': tfds.features.Tensor(
                        shape=(7,),
                        dtype=np.float32,
                        doc='Action consists of [3x robot delta displacements, 3x delta rotations, and gripper closedness action].',
                    ),
                    'discount': tfds.features.Scalar(
                        dtype=np.float32,
                        doc='Discount if provided, default to 1.'
                    ),
                    'reward': tfds.features.Scalar(
                        dtype=np.float32,
                        doc='Reward if provided, 1 on final step for demos.'
                    ),
                    'is_first': tfds.features.Scalar(
                        dtype=np.bool_,
                        doc='True on the first step of the episode.'
                    ),
                    'is_last': tfds.features.Scalar(
                        dtype=np.bool_,
                        doc='True on the last step of the episode.'
                    ),
                    'is_terminal': tfds.features.Scalar(
                        dtype=np.bool_,
                        doc='True on the last step of the episode if it is a terminal step, True for demos.'
                    ),
                    'natural_language_instruction': tfds.features.Text(
                        doc='Language Instruction.'
                    ),
                    'natural_language_embedding': tfds.features.Tensor(
                        shape=(512,),
                        dtype=np.float32,
                        doc='Kona language embedding. See https://tfhub.dev/google/universal-sentence-encoder-large/5'
                    ),
                }),
                'episode_metadata': tfds.features.FeaturesDict({
                    'file_path': tfds.features.Text(
                        doc='Path to the original data file.'
                    ),
                }),
            }))

    def _split_generators(self, dl_manager: tfds.download.DownloadManager):
        """Define data splits."""
        return {
            'train': self._generate_examples(path='data/train/episode_*.npy'),
        }

    def _generate_examples(self, path) -> Iterator[Tuple[str, Any]]:
        """Generator of examples for each split."""

        def _parse_example(episode_path):
            # Load raw data --> this should change for your dataset
            data = np.load(episode_path, allow_pickle=True).item()

            # Extract data
            images_front = data["images_front"]
            images_side = data["images_side"]
            images_hand = data["images_hand"]
            states = data["states"]
            actions = data["actions"]
            natural_language_instruction = data["natural_language"]

            # Verify that the lengths of the data match
            if not (len(images_front) == len(images_side) == len(images_hand) == len(states) == len(actions)):
                raise ValueError(f"Mismatch in lengths of images, states, and actions in `{episode_path}`")

            # Assemble episode
            episode = []
            for i in range(len(states)):
                # Compute language embedding (use a fixed instruction for now)
                language_embedding = self._embed([natural_language_instruction[0]])[0].numpy()

                episode.append({
                    'observation': {
                        'image_front': images_front[i],
                        'image_side': images_side[i],
                        'image_hand': images_hand[i],
                        'state': states[i],
                    },
                    'action': actions[i],
                    'discount': 1.0,
                    'reward': float(i == (len(states) - 1)),
                    'is_first': i == 0,
                    'is_last': i == (len(states) - 1),
                    'is_terminal': i == (len(states) - 1),
                    'natural_language_instruction': natural_language_instruction[i],
                    'natural_language_embedding': language_embedding,
                })

            # Create output data sample
            sample = {
                'steps': episode,
                'episode_metadata': {
                    'file_path': episode_path
                }
            }

            return episode_path, sample

        # Create a list of all examples
        episode_paths = glob.glob(path)

        # For small datasets, use single-thread parsing
        for sample in episode_paths:
            yield _parse_example(sample)

The generated datasets are then saved on google drive and uploaded in the Colab script:

python

# Install the gdown library if not already present
!pip install gdown --quiet

# Specify the Google Drive file ID
file_id = "<Your_shared_Google_Drive_link>"

# Download the file from Google Drive
!gdown $file_id -O /content/overall900.zip

# Load the dataset
import os

# Path to the ZIP file
zip_path_data = "/content/overall900.zip"

# Path to the output directory
output_dir = "/content/openvla/overall900_dataset"

# Create the output directory if it doesn't exist
os.makedirs(output_dir, exist_ok=True)

# Extract the ZIP file to the specified directory
!unzip -o "$zip_path_data" -d "$output_dir"

The finetuning script

Prior to run this script remember to open an wandb account:

python

!pip install torch==2.2.0 accelerate==0.24.1 timm==0.9.10 peft==0.11.1 wandb==0.18.7 draccus==0.8.0 transformers==4.40.1

# Import libraries
import os
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
import timm
import torch
import torch.nn as nn
from typing import Callable
from PIL import Image
import torch.distributed as dist
import tqdm
from accelerate import PartialState
from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForVision2Seq, AutoProcessor, BitsAndBytesConfig
from transformers import AutoConfig, AutoImageProcessor
from transformers.modeling_outputs import CausalLMOutputWithPast
import wandb
from prismatic.models.backbones.llm.prompting import PurePromptBuilder, VicunaV15ChatPromptBuilder
from prismatic.util.data_utils import PaddedCollatorForActionPrediction
from prismatic.vla.action_tokenizer import ActionTokenizer
from prismatic.vla.datasets import RLDSBatchTransform, RLDSDataset
from prismatic.vla.datasets.rlds.utils.data_utils import save_dataset_statistics
from prismatic.extern.hf.configuration_prismatic import OpenVLAConfig
from prismatic.extern.hf.modeling_prismatic import OpenVLAForActionPrediction
from prismatic.extern.hf.processing_prismatic import PrismaticImageProcessor, PrismaticProcessor

# Log in to Weights & Biases
wandb.login(key="YOUR_WANDB_KEY")  # Keep your key secure!

# Disable tokenizer parallelism to avoid conflicts on a single GPU
os.environ["TOKENIZERS_PARALLELISM"] = "false"

@dataclass
class FinetuneConfig:
    vla_path: str = "openvla/openvla-7b"  # Path to the OpenVLA model
    data_root_dir: Path = Path("/content/openvla/overall900_dataset")  # Dataset path
    dataset_name: str = "overall900"  # Dataset name
    run_root_dir: Path = Path("checkpoints")  # Path for logs and checkpoints
    adapter_tmp_dir: Path = Path("adapter-tmp")  # Temporary directory for LoRA weights

    # Training parameters
    batch_size: int = 16
    max_steps: int = 1200
    save_steps: int = 1200
    learning_rate: float = 5e-4
    grad_accumulation_steps: int = 2
    image_aug: bool = False
    shuffle_buffer_size: int = 600
    save_latest_checkpoint_only: bool = True

    # LoRA and quantization settings
    use_lora: bool = True
    lora_rank: int = 32
    lora_dropout: float = 0.1
    use_quantization: bool = True

    # Monitoring parameters
    wandb_project: str = "model-overall900"
    wandb_entity: str = "YOUR_WANDB_ENTITY"
    run_id_note: Optional[str] = None

# Finetuning function
def finetune(cfg: FinetuneConfig) -> None:
    print(f"Fine-tuning OpenVLA Model `{cfg.vla_path}` on `{cfg.dataset_name}`")

    # Validate GPU availability
    assert torch.cuda.is_available(), "Fine-tuning requires at least one GPU!"
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)
    torch.cuda.empty_cache()

    # Configure experiment ID
    exp_id = (
        f"{cfg.vla_path.split('/')[-1]}+{cfg.dataset_name}"
        f"+b{cfg.batch_size * cfg.grad_accumulation_steps}"
        f"+lr-{cfg.learning_rate}"
    )
    if cfg.use_lora:
        exp_id += f"+lora-r{cfg.lora_rank}+dropout-{cfg.lora_dropout}"
    if cfg.use_quantization:
        exp_id += "+q-4bit"
    if cfg.run_id_note:
        exp_id += f"--{cfg.run_id_note}"
    if cfg.image_aug:
        exp_id += "--image_aug"

    run_dir, adapter_dir = cfg.run_root_dir / exp_id, cfg.adapter_tmp_dir / exp_id
    os.makedirs(run_dir, exist_ok=True)

    quantization_config = None
    if cfg.use_quantization:
        assert cfg.use_lora, "Quantization is only supported for LoRA fine-tuning!"
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_type="nf4"
        )

    # Register the model with HuggingFace
    AutoConfig.register("openvla", OpenVLAConfig)
    AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor)
    AutoProcessor.register(OpenVLAConfig, PrismaticProcessor)
    AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction)

    processor = AutoProcessor.from_pretrained(cfg.vla_path, trust_remote_code=True)
    vla = AutoModelForVision2Seq.from_pretrained(
        cfg.vla_path,
        torch_dtype=torch.bfloat16,
        quantization_config=quantization_config,
        attn_implementation="flash_attention_2",
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )

    if cfg.use_quantization:
        vla = prepare_model_for_kbit_training(vla)
    else:
        vla = vla.to(device)

    if cfg.use_lora:
        lora_config = LoraConfig(
            r=cfg.lora_rank,
            lora_alpha=min(cfg.lora_rank, 16),
            lora_dropout=cfg.lora_dropout,
            target_modules="all-linear",
            init_lora_weights="gaussian",
        )
        vla = get_peft_model(vla, lora_config)
        vla.print_trainable_parameters()

    vla.to(device)

    # Create optimizer
    trainable_params = [param for param in vla.parameters() if param.requires_grad]
    optimizer = AdamW(trainable_params, lr=cfg.learning_rate)
    action_tokenizer = ActionTokenizer(processor.tokenizer)

    batch_transform = RLDSBatchTransform(
        action_tokenizer,
        processor.tokenizer,
        image_transform=processor.image_processor.apply_transform,
        prompt_builder_fn=PurePromptBuilder if "v01" not in cfg.vla_path else VicunaV15ChatPromptBuilder,
    )

    vla_dataset = RLDSDataset(
        cfg.data_root_dir,
        cfg.dataset_name,
        batch_transform,
        resize_resolution=tuple(vla.config.image_sizes),
        shuffle_buffer_size=cfg.shuffle_buffer_size,
        image_aug=cfg.image_aug,
    )

    # Save dataset statistics
    save_dataset_statistics(vla_dataset.dataset_statistics, run_dir)

    collator = PaddedCollatorForActionPrediction(
        processor.tokenizer.model_max_length, processor.tokenizer.pad_token_id, padding_side="right"
    )

    dataloader = DataLoader(
        vla_dataset,
        batch_size=cfg.batch_size,
        collate_fn=collator,
        num_workers=0,
    )

    # Initialize Weights & Biases logging
    wandb.init(entity=cfg.wandb_entity, project=cfg.wandb_project, name=f"ft+{exp_id}", force=True)

    recent_losses = deque(maxlen=cfg.grad_accumulation_steps)
    recent_action_accuracies = deque(maxlen=cfg.grad_accumulation_steps)
    recent_l1_losses = deque(maxlen=cfg.grad_accumulation_steps)

    # Training loop
    with tqdm.tqdm(total=cfg.max_steps, leave=False) as progress:
        vla.train()
        optimizer.zero_grad()
        for batch_idx, batch in enumerate(dataloader):
            with torch.autocast("cuda", dtype=torch.bfloat16):
                output: CausalLMOutputWithPast = vla(
                    input_ids=batch["input_ids"].to(device),
                    attention_mask=batch["attention_mask"].to(device),
                    pixel_values=batch["pixel_values"].to(torch.bfloat16).to(device),
                    labels=batch["labels"],
                )
                loss = output.loss

            # Normalize loss for gradient accumulation
            normalized_loss = loss / cfg.grad_accumulation_steps
            normalized_loss.backward()

            # Compute accuracy and L1 loss for logging
            action_logits = output.logits[:, vla.vision_backbone.featurizer.patch_embed.num_patches : -1]
            action_preds = action_logits.argmax(dim=2).to(device)
            action_gt = batch["labels"][:, 1:].to(action_preds.device).to(device)
            mask = action_gt > action_tokenizer.action_token_begin_idx

            # Compute accuracy
            correct_preds = (action_preds == action_gt) & mask
            action_accuracy = correct_preds.sum().float() / mask.sum().float()

            # Compute L1 loss on predicted (continuous) actions
            continuous_actions_pred = torch.tensor(
                action_tokenizer.decode_token_ids_to_actions(action_preds[mask].cpu().numpy())
            ).to(device)
            continuous_actions_gt = torch.tensor(
                action_tokenizer.decode_token_ids_to_actions(action_gt[mask].cpu().numpy())
            ).to(device)
            action_l1_loss = torch.nn.functional.l1_loss(continuous_actions_pred, continuous_actions_gt)

            # Store recent metrics
            recent_losses.append(loss.item())
            recent_action_accuracies.append(action_accuracy.item())
            recent_l1_losses.append(action_l1_loss.item())

            # Compute gradient step index
            gradient_step_idx = batch_idx // cfg.grad_accumulation_steps

            # Log metrics to Weights & Biases every 10 gradient steps
            if gradient_step_idx % 10 == 0:
                wandb.log(
                    {
                        "train_loss": sum(recent_losses) / len(recent_losses),
                        "action_accuracy": sum(recent_action_accuracies) / len(recent_action_accuracies),
                        "l1_loss": sum(recent_l1_losses) / len(recent_l1_losses),
                    },
                    step=gradient_step_idx,
                )

            # Optimizer step
            if (batch_idx + 1) % cfg.grad_accumulation_steps == 0:
                optimizer.step()
                optimizer.zero_grad()
                progress.update()

            # Save model checkpoint
            if gradient_step_idx > 0 and gradient_step_idx % cfg.save_steps == 0:
                print(f"Saving Model Checkpoint for Step {gradient_step_idx}")
                save_dir = adapter_dir if cfg.use_lora else run_dir
                processor.save_pretrained(run_dir)
                vla.save_pretrained(save_dir)

                if cfg.use_lora:
                    base_vla = AutoModelForVision2Seq.from_pretrained(
                        cfg.vla_path, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, trust_remote_code=True
                    )
                    merged_vla = PeftModel.from_pretrained(base_vla, adapter_dir)
                    merged_vla = merged_vla.merge_and_unload()
                    if cfg.save_latest_checkpoint_only:
                        merged_vla.save_pretrained(run_dir)
                        print(f"Saved Model Checkpoint for Step {gradient_step_idx} at: {run_dir}")
                    else:
                        checkpoint_dir = Path(str(run_dir) + f"--{gradient_step_idx}_chkpt")
                        os.makedirs(checkpoint_dir, exist_ok=True)
                        save_dataset_statistics(vla_dataset.dataset_statistics, checkpoint_dir)
                        processor.save_pretrained(checkpoint_dir)
                        merged_vla.save_pretrained(checkpoint_dir)
                        print(f"Saved Model Checkpoint for Step {gradient_step_idx} at: {checkpoint_dir}")

            # Stop training when max_steps is reached
            if gradient_step_idx == cfg.max_steps:
                print(f"Max step {cfg.max_steps} reached! Stopping training...")
                break

# Start fine-tuning in Colab
if __name__ == "__main__":
    cfg = FinetuneConfig()
    finetune(cfg)

Saving checkpoints

Checkpoints are then being saved by compressing the checkpoint folder into a .tar file and then uploading it directly in Google drive:

python

import shutil

# Path to the folder to compress
checkpoint_folder = '/content/openvla/checkpoints'

# Name of the .tar file
tar_filename = '/content/checkpoints_backup.tar'

# Create the .tar archive
shutil.make_archive(tar_filename.replace('.tar', ''), 'tar', checkpoint_folder)

# Mount Google Drive
from google.colab import drive
drive.mount('/content/drive')

# Copy the .tar file to Google Drive
!cp /content/checkpoints_backup.tar /content/drive/MyDrive/

# Sync to ensure the file is written to Google Drive
!sync

merlo00/OpenVLA-IsaacLab-finetuning

Google Colab Finetuning Pipeline for OpenVLA

Table of Contents

Introduction

OpenVLA folder editing

Package installation

Dataset loading

The finetuning script

Saving checkpoints

merlo00/
OpenVLA-IsaacLab-finetuning