diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json new file mode 100644 index 000000000..a8eb7a295 --- /dev/null +++ b/.devcontainer/devcontainer.json @@ -0,0 +1,36 @@ +{ + "name": "OpenVLA Development", + "dockerComposeFile": "docker-compose.yml", + "service": "devcontainer", + "workspaceFolder": "/workspace/openvla", + "containerEnv": { + "PYTHONPATH": "${containerWorkspaceFolder}" + }, + "remoteEnv": { + // Environment variables will be loaded from .env + }, + "customizations": { + "vscode": { + "extensions": [ + "ms-python.python", + "ms-python.vscode-pylance", + "ms-toolsai.jupyter", + "github.copilot" + ], + "settings": { + "python.defaultInterpreterPath": "/usr/local/bin/python", + "python.linting.enabled": true, + "python.linting.pylintEnabled": true, + "python.formatting.provider": "black", + "editor.formatOnSave": true, + "editor.rulers": [121] + } + } + }, + "remoteUser": "root", + "postCreateCommand": "pip install -e .", + // Load environment variables from .env file + "features": { + "ghcr.io/devcontainers/features/dotnet:1": {} + } + } \ No newline at end of file diff --git a/.devcontainer/docker-compose.yml b/.devcontainer/docker-compose.yml new file mode 100644 index 000000000..453872554 --- /dev/null +++ b/.devcontainer/docker-compose.yml @@ -0,0 +1,28 @@ +services: + devcontainer: + build: + context: .. + dockerfile: ./Dockerfile + working_dir: /workspace/openvla + command: sleep infinity + volumes: + - ..:/workspace/openvla:cached + - finetuner-cache:/root/.cache + - ${HOME}/.cache/huggingface:/root/.cache/huggingface # use cached models/datasets from host + environment: + - WANDB_API_KEY=${WANDB_API_KEY} + - WANDB_MODE=online + - HF_HOME=/root/.cache/huggingface + - HF_TOKEN=${HF_TOKEN} + shm_size: 16gb + network_mode: host + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + +volumes: + finetuner-cache: \ No newline at end of file diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 000000000..fdf1b4af0 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,13 @@ +# Prevent these from being sent to Docker daemon +.git/ +.venv/ +__pycache__/ +*.py[cod] +.DS_Store +.env +data/ +.runs/ +.adapter/ +*.swp +*.swo +*.log \ No newline at end of file diff --git a/.env.template b/.env.template new file mode 100644 index 000000000..30f2e5d00 --- /dev/null +++ b/.env.template @@ -0,0 +1,3 @@ +HF_TOKEN= +WANDB_API_KEY= +WANDB_PROJECT="ur5e" \ No newline at end of file diff --git a/.gitignore b/.gitignore index 6160ebcb1..4dd11c4ae 100644 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,6 @@ data/ # Rollout videos and wandb logs rollouts/ wandb/ +.tmp +.runs +.preprocessors diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 000000000..bce919cf1 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/lerobot"] + path = third_party/lerobot + url = https://github.com/huggingface/lerobot.git diff --git a/ALOHA.md b/ALOHA.md new file mode 100644 index 000000000..bc508fcfe --- /dev/null +++ b/ALOHA.md @@ -0,0 +1,158 @@ +# OpenVLA-OFT+ in Real-World ALOHA Robot Tasks + +## Relevant Files + +Evaluation +* `experiments/robot/aloha/`: ALOHA training and eval files + * `run_aloha_eval.py`: ALOHA eval script (CLIENT SIDE; see "SERVER SIDE" below) + * `aloha_utils.py`: ALOHA eval utils + * Other ALOHA robot environment files copied from the original [ALOHA GitHub repo](https://github.com/tonyzhaozh/aloha): + * `constants.py` + * `real_env.py` + * `robot_utils.py` +* `experiments/robot/`: General eval utils files + * `openvla_utils.py`: OpenVLA-specific eval utils + * `robot_utils.py`: Other eval utils +* `vla-scripts/deploy.py`: VLA server deploy script (SERVER SIDE) + +Note: Unlike the LIBERO evaluation setup, we use a server-client interface here. This is particularly useful if the user's machine which commands the robot does not have access to a local GPU with sufficient specs to run the fine-tuned VLA policies. + +Training +* `experiments/robot/aloha/`: ALOHA training and eval files + * `preprocess_split_aloha_data.py`: ALOHA data preprocessing script +* `vla-scripts/finetune.py`: VLA fine-tuning script + +## Setup + +Set up a conda environment for training policies and deploying them on the VLA server (see instructions in [SETUP.md](SETUP.md)). + +## Fine-Tuning on ALOHA Robot Data + +We assume that you have collected a set of expert demonstrations on the ALOHA robot already. + +First, use our `preprocess_split_aloha_data.py` script to preprocess the raw ALOHA dataset: downsize images from 480x640 to 256x256 and split into training and validation sets. Below are examples for the `put X into pot` task in our paper (which has 3 possible target objects, 1 per episode): + +```bash +python experiments/robot/aloha/preprocess_split_aloha_data.py \ + --dataset_path /scr/moojink/data/aloha1_raw/put_green_pepper_into_pot/ \ + --out_base_dir /scr/moojink/data/aloha1_preprocessed/ \ + --percent_val 0.05 +python experiments/robot/aloha/preprocess_split_aloha_data.py \ + --dataset_path /scr/moojink/data/aloha1_raw/put_red_pepper_into_pot/ \ + --out_base_dir /scr/moojink/data/aloha1_preprocessed/ \ + --percent_val 0.05 +python experiments/robot/aloha/preprocess_split_aloha_data.py \ + --dataset_path /scr/moojink/data/aloha1_raw/put_yellow_corn_into_pot/ \ + --out_base_dir /scr/moojink/data/aloha1_preprocessed/ \ + --percent_val 0.05 +``` + +Then, convert the preprocessed ALOHA datasets into a single RLDS dataset that is compatible with OpenVLA fine-tuning. This process is the same as in the original OpenVLA repo. See instructions for converting to RLDS [here](https://github.com/moojink/rlds_dataset_builder) (a sample ALOHA preprocessed-to-RLDS conversion script is available [here](https://github.com/moojink/rlds_dataset_builder/blob/main/aloha1_put_X_into_pot_300_demos/aloha1_put_X_into_pot_300_demos_dataset_builder.py); this script converts the three preprocessed datasets above into one unified RLDS dataset, with train/val splits). + +After converting to RLDS, register the dataset (which, for the example task above, would be called `aloha1_put_X_into_pot_300_demos`) with our dataloader by adding an entry for it in `configs.py` ([here](prismatic/vla/datasets/rlds/oxe/configs.py#L680)), `transforms.py` ([here](prismatic/vla/datasets/rlds/oxe/transforms.py#L928)), and `mixtures.py` ([here](prismatic/vla/datasets/rlds/oxe/mixtures.py#L216)). For reference, in each of these files, there are sample entries for the ALOHA datasets that we used in our paper. + +Before fine-tuning, set the desired ALOHA action chunk size in [`prismatic/vla/constants.py`](prismatic/vla/constants.py) (see `NUM_ACTIONS_CHUNK` in `ALOHA_CONSTANTS`). We set it to 25 by default because we used a control frequency of 25 Hz in our ALOHA setup to reduce storage costs and training time (while still maintaining smoothness in the robot's motions). If you use 50 Hz, we recommend setting `NUM_ACTIONS_CHUNK` to `50`. In general, 1 second-long action chunks are a good default. Do NOT modify `ACTION_PROPRIO_NORMALIZATION_TYPE`: Since the ALOHA robot action space is absolute joint angles, we do not want to use a normalization scheme that clips outlier values (like the Q1-Q99 normalization we used with the relative end-effector pose actions for LIBERO), since that would prevent the model from outputting certain robot joint angles that are crucial for solving the task. + +Now begin fine-tuning! Below is a sample command to fine-tune OpenVLA using our OFT+ recipe on the `put X into pot` task above ("+" in "OFT+" means FiLM is included for enhanced language grounding). Replace `X` in the first line with the number of GPUs available to you. + +```bash +torchrun --standalone --nnodes 1 --nproc-per-node X vla-scripts/finetune.py \ + --vla_path openvla/openvla-7b \ + --data_root_dir /PATH/TO/RLDS/DATASETS/DIR/ \ + --dataset_name aloha1_put_X_into_pot_300_demos \ + --run_root_dir /YOUR/CHECKPOINTS/AND/LOG/DIR/ \ + --use_l1_regression True \ + --use_diffusion False \ + --use_film True \ + --num_images_in_input 3 \ + --use_proprio True \ + --batch_size 4 \ + --learning_rate 5e-4 \ + --num_steps_before_decay 50000 \ + --max_steps 100005 \ + --use_val_set True \ + --val_freq 10000 \ + --save_freq 10000 \ + --save_latest_checkpoint_only False \ + --image_aug True \ + --lora_rank 32 \ + --wandb_entity "YOUR_WANDB_ENTITY" \ + --wandb_project "YOUR_WANDB_PROJECT" \ + --run_id_note parallel_dec--25_acts_chunk--continuous_acts--L1_regression--3rd_person_img--left_right_wrist_imgs--proprio_state--film +``` + +The above training command should reproduce our OpenVLA-OFT+ results on the `put X into pot` task if `X = 8` and the 100K step checkpoint is evaluated. It will fine-tune OpenVLA using 3 input images (1 third-person image + 2 wrist camera images). Note that we use learning rate decay after a certain point (50K steps in the command above) since doing so speeds up training convergence (train L1 loss spikes down from our experience). + +Best practices for fine-tuning: +* In general, we recommend fine-tuning until training L1 loss goes below 0.01 and starts to plateau. + * One way to achieve this is to fine-tune using our default learning rate of `5e-4` until the loss starts to decrease very slowly, and then decay the learning rate by 10x to `5e-5` (which should make the loss spike down) and train until the training L1 loss finally plateaus. +* Depending on your dataset size, you may need to adjust some hyperparameters. For example, if you use a large dataset with over 300 demos, you may need to decay the learning rate later and train for longer for best performance. Decaying too earlier can lead to a suboptimal policy. +* If your task does not require good langauge grounding (e.g., if there is only one language instruction), FiLM is not necessary; consider setting `--use_film False` to train fewer model parameters. +* Please be sure to test your policy with the same device/GPU used to train it! Otherwise, performance may drop substantially. You may be able to avoid the performance drop if you merge the LoRA weights into the base model on the downstream device used for testing (e.g., if you train on H100 and then merge on A100 before testing on A100). You can see our script [vla-scripts/merge_lora_weights_and_save.py](vla-scripts/merge_lora_weights_and_save.py) for merging the LoRA adapter into the base model offline. It's okay if you already merged LoRA weights into the base OpenVLA model during fine-tuning; you can always redownload the base model and merge again as long as you still have the LoRA adapter (`merge_lora_weights_and_save.py` will handle this for you). + +If you run into any issues, please open a new GitHub issue. + +## Launching ALOHA Robot Evaluations + +In the primary conda environment (`openvla-oft`) which you will use to launch the VLA server, install a few packages for the server-client interface: + +```bash +conda activate openvla-oft +pip install uvicorn fastapi json-numpy +``` + +On the machine that you will use to command the robot, set up a second conda environment that will be used to run the robot environment, query the VLA server, and execute actions in the environment: + +```bash +# Create and activate client conda environment +conda create -n openvla-oft-aloha python=3.10 -y +conda activate openvla-oft-aloha + +# 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 packages needed for the ALOHA robot environment +pip install -r experiments/robot/aloha/requirements_aloha.txt +``` + +Launch the VLA server on the machine that has the GPU you will use to run model inference (using the `openvla-oft` conda environment). Below is a sample command for this (change as needed): + +```bash +python vla-scripts/deploy.py \ + --pretrained_checkpoint /PATH/TO/FINETUNED/MODEL/CHECKPOINT/DIR/ \ + --use_l1_regression True \ + --use_film True \ + --num_images_in_input 3 \ + --use_proprio True \ + --center_crop True \ + --num_open_loop_steps 25 \ + --unnorm_key aloha1_put_X_into_pot_300_demos +``` + +Then, run the ALOHA evaluation script. Specify the VLA server URL or IP address in the `vla_server_url` argument. Below is a sample command: + +```bash +python experiments/robot/aloha/run_aloha_eval.py \ + --center_crop True \ + --num_open_loop_steps 25 \ + --use_vla_server True \ + --vla_server_url \ + --num_rollouts_planned \ + --max_steps +``` + +If you run into any issues, please open a new GitHub issue. + +## Troubleshooting Tips + +* Tip #1: If you run into a ROS error such as `ImportError: /lib/x86_64-linux-gnu/libp11-kit.so.0: undefined symbol: ffi_type_pointer, version LIBFFI_BASE_7.0`, try running the following command in your client conda environment (`openvla-oft-aloha`): + + ``` + conda install -c conda-forge libffi + ``` diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000..82e569c00 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,31 @@ +FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-devel + +# Install system dependencies for flash-attn. +RUN apt-get update && apt-get install -y \ + git \ + ninja-build \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace +RUN git clone -b merge-finetuner-changes https://github.com/nomagiclab/openvla.git && \ + cd openvla && \ + git submodule init third_party/lerobot && \ + git submodule update --recursive --init third_party/lerobot + +WORKDIR /workspace/openvla + +RUN ls -la third_party/lerobot + +# Editable install of openvla, then lerobot submodule, then reinstall newly +# missing openvla dependencies to negotiate dependency incompatibility. +# Then install flash-attn separately (per OpenVLA instructions) +# and download the openvla-7b model. +RUN pip install -e . && \ + cd third_party/lerobot/ && \ + pip install -e . && \ + cd ../../ && \ + pip check | awk '$1 ~ /openvla/ {gsub(/,/,"",$5); print $5}' | xargs pip install && \ + pip install packaging ninja && \ + pip install "flash-attn==2.5.5" --no-build-isolation && \ + pip install huggingface-hub && \ + huggingface-cli download openvla/openvla-7b diff --git a/LIBERO.md b/LIBERO.md new file mode 100644 index 000000000..851af31cd --- /dev/null +++ b/LIBERO.md @@ -0,0 +1,123 @@ +# OpenVLA-OFT in the LIBERO Simulation Benchmark + +## Relevant Files + +Evaluation +* `experiments/robot/libero/`: LIBERO eval files + * `run_libero_eval.py`: LIBERO eval script + * `libero_utils.py`: LIBERO eval utils +* `experiments/robot/`: General eval utils files + * `openvla_utils.py`: OpenVLA-specific eval utils + * `robot_utils.py`: Other eval utils + +Training +* `vla-scripts/finetune.py`: VLA fine-tuning script + + +## Setup + +Set up a conda environment (see instructions in [SETUP.md](SETUP.md)). + +Clone and install the [LIBERO repo](https://github.com/Lifelong-Robot-Learning/LIBERO) and required packages: + +```bash +git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git +pip install -e LIBERO +pip install -r experiments/robot/libero/libero_requirements.txt # From openvla-oft base dir +``` + +(Optional, if you plan to launch training) To download the [LIBERO datasets](https://huggingface.co/datasets/openvla/modified_libero_rlds) that we used in our fine-tuning +experiments, run the command below. This will download the LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, +and LIBERO-10 datasets in RLDS data format (~10 GB total). You can use these to fine-tune OpenVLA or +train other methods. This step is optional since we provide pretrained OpenVLA-OFT checkpoints below. +Note that these are the same datasets used in the original OpenVLA project. If needed, see details on how to download the original non-RLDS datasets [here](https://github.com/openvla/openvla?tab=readme-ov-file#libero-setup). +```bash +git lfs clone git@hf.co:datasets/openvla/modified_libero_rlds +``` + +## Launching LIBERO Evaluations + +We fine-tuned OpenVLA via LoRA (r=32) with our OFT recipe on four LIBERO task suites independently: LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-10 (also called LIBERO-Long). +The four OpenVLA-OFT checkpoints for LIBERO are available on Hugging Face: +* [moojink/openvla-7b-oft-finetuned-libero-spatial](https://huggingface.co/moojink/openvla-7b-oft-finetuned-libero-spatial) +* [moojink/openvla-7b-oft-finetuned-libero-object](https://huggingface.co/moojink/openvla-7b-oft-finetuned-libero-object) +* [moojink/openvla-7b-oft-finetuned-libero-goal](https://huggingface.co/moojink/openvla-7b-oft-finetuned-libero-goal) +* [moojink/openvla-7b-oft-finetuned-libero-10](https://huggingface.co/moojink/openvla-7b-oft-finetuned-libero-10) + +To start evaluations with one of these checkpoints, run one of the commands below. Each will automatically download the appropriate checkpoint listed above. You can set the `TRANSFORMERS_CACHE` and `HF_HOME` environment variable to change where the checkpoint files get cached. + +```bash +# Launch LIBERO-Spatial evals +python experiments/robot/libero/run_libero_eval.py \ + --pretrained_checkpoint moojink/openvla-7b-oft-finetuned-libero-spatial \ + --task_suite_name libero_spatial + +# Launch LIBERO-Object evals +python experiments/robot/libero/run_libero_eval.py \ + --pretrained_checkpoint moojink/openvla-7b-oft-finetuned-libero-object \ + --task_suite_name libero_object + +# Launch LIBERO-Goal evals +python experiments/robot/libero/run_libero_eval.py \ + --pretrained_checkpoint moojink/openvla-7b-oft-finetuned-libero-goal \ + --task_suite_name libero_goal + +# Launch LIBERO-10 (LIBERO-Long) evals +python experiments/robot/libero/run_libero_eval.py \ + --pretrained_checkpoint moojink/openvla-7b-oft-finetuned-libero-10 \ + --task_suite_name libero_10 +``` + +Notes: +* The evaluation script will run 500 trials by default (10 tasks x 50 episodes each). You can modify the number of + trials per task by setting `--num_trials_per_task`. You can also change the random seed via `--seed`. There are + other arguments in the script; we set them to the default values that work with the OpenVLA-OFT checkpoints above. +* **NOTE: Setting `--center_crop True` is important** because we fine-tuned OpenVLA with random crop augmentations + (we took a random crop with 90% area in every training sample, so at test time we simply take the center 90% crop). +* The evaluation script logs results locally. You can also log results in Weights & Biases + by setting `--use_wandb True` and specifying `--wandb_project ` and `--wandb_entity `. +* The results reported in our paper were obtained using **Python 3.10.14, PyTorch 2.2.0, and our + [custom transformers v4.40.1 fork](https://github.com/moojink/transformers-openvla-oft.git)** + on an **NVIDIA A100 GPU**, averaged over three random seeds. Please stick to these package versions if possible. + Note that results may vary slightly if you use a different GPU than the A100. If the discrepancy is large, + please post a GitHub issue, and we will look into it. + +## Fine-Tuning on LIBERO Datasets + +First, download the LIBERO datasets as mentioned above in the Setup section above: `libero_spatial_no_noops`, `libero_object_no_noops`, `libero_goal_no_noops`, `libero_10_no_noops`. (`"_no_noops"` stands for no no-op actions, i.e., training samples with near-zero actions are filtered out). + +Then, launch the fine-tuning script with the OFT configuration below, replacing `X` in the first line with the number of GPUs. The command below launches fine-tuning on LIBERO-Spatial with the hyperparameters that we used in our paper. Here, batch size 8 per GPU will require ~62 GB VRAM, and batch size 1 per GPU will require ~25 GB VRAM. + +```bash +torchrun --standalone --nnodes 1 --nproc-per-node X vla-scripts/finetune.py \ + --vla_path openvla/openvla-7b \ + --data_root_dir /PATH/TO/RLDS/DATASETS/DIR/ \ + --dataset_name libero_spatial_no_noops \ + --run_root_dir /YOUR/CHECKPOINTS/AND/LOG/DIR/ \ + --use_l1_regression True \ + --use_diffusion False \ + --use_film False \ + --num_images_in_input 2 \ + --use_proprio True \ + --batch_size 8 \ + --learning_rate 5e-4 \ + --num_steps_before_decay 100000 \ + --max_steps 150005 \ + --save_freq 10000 \ + --save_latest_checkpoint_only False \ + --image_aug True \ + --lora_rank 32 \ + --wandb_entity "YOUR_WANDB_ENTITY" \ + --wandb_project "YOUR_WANDB_PROJECT" \ + --run_id_note parallel_dec--8_acts_chunk--continuous_acts--L1_regression--3rd_person_img--wrist_img--proprio_state +``` + +The above training command should reproduce our OpenVLA-OFT results if `X = 8` and the 150K step checkpoint is evaluated. + +You can replace `libero_spatial_no_noops` with `libero_object_no_noops`, `libero_goal_no_noops`, or `libero_10_no_noops`. You can also modify other args — e.g., if you want to train with just one input image from the third-person camera and disable proprio state input, you can set `--num_images_in_input 1` and `--use_proprio False`. + +In general, we recommend fine-tuning until training L1 loss goes below 0.01 and starts to plateau (with the above configuration, it should reach ~0.006 L1 loss on LIBERO-Spatial after 150K gradient steps with 10x LR decay after 100K steps). However, for LIBERO-Goal only, we found that the 50K checkpoint (which was at ~0.02 L1 loss) performed best for unknown reasons. For all other task suites though, we found that the 150K checkpoint performed best. + +Please be sure to test your policy with the same device/GPU used to train it! Otherwise, performance may drop substantially. You may be able to avoid the performance drop if you merge the LoRA weights into the base model on the downstream device used for testing (e.g., if you train on H100 and then merge on A100 before testing on A100). You can see our script [vla-scripts/merge_lora_weights_and_save.py](vla-scripts/merge_lora_weights_and_save.py) for merging the LoRA adapter into the base model offline. It's okay if you already merged LoRA weights into the base OpenVLA model during fine-tuning; you can always redownload the base model and merge again as long as you still have the LoRA adapter (`merge_lora_weights_and_save.py` will handle this for you). + +If you run into any issues, please open a new GitHub issue. If you do not receive a response within 2 business days, please email Moo Jin Kim (moojink@cs.stanford.edu) to bring the issue to his attention. diff --git a/LICENSE b/LICENSE index 04f26a7d0..b2c22d519 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2024 Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti. +Copyright (c) 2025 Moo Jin Kim, Chelsea Finn, Percy Liang. Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/README.md b/README.md index ef62188ed..2ee12646a 100644 --- a/README.md +++ b/README.md @@ -1,646 +1,97 @@ -# OpenVLA: An Open-Source Vision-Language-Action Model +# Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success -[![arXiv](https://img.shields.io/badge/arXiv-2406.09246-df2a2a.svg?style=for-the-badge)](https://arxiv.org/abs/2406.09246) -[![HF Models](https://img.shields.io/badge/%F0%9F%A4%97-Models-yellow?style=for-the-badge)](https://huggingface.co/openvla/openvla-7b) -[![PyTorch](https://img.shields.io/badge/PyTorch-2.2.0-EE4C2C.svg?style=for-the-badge&logo=pytorch)](https://pytorch.org/get-started/locally/) -[![Python](https://img.shields.io/badge/python-3.10-blue?style=for-the-badge)](https://www.python.org) -[![License](https://img.shields.io/github/license/TRI-ML/prismatic-vlms?style=for-the-badge)](LICENSE) - -[**Getting Started**](#getting-started) | [**Pretrained VLAs**](#pretrained-vlas) | [**Installation**](#installation) | [**Fine-Tuning OpenVLA via LoRA**](#fine-tuning-openvla-via-lora) | [**Fully Fine-Tuning OpenVLA**](#fully-fine-tuning-openvla) | -[**Training VLAs from Scratch**](#training-vlas-from-scratch) | [**Evaluating OpenVLA**](#evaluating-openvla) | [**Project Website**](https://openvla.github.io/) +**Project website: https://openvla-oft.github.io/** +**Paper: https://arxiv.org/abs/2502.19645** -
+**Summary video: https://youtu.be/T3Zkkr_NTSA** -## Latest Updates -- [2025-03-03] OFT (Optimized Fine-Tuning recipe for VLAs) was recently released! Compared to vanilla OpenVLA fine-tuning, OFT enables 25-50x faster inference speed, higher task success rates, multiple input images, and high-frequency bimanual robot control. Unlike FAST, OFT uses continuous actions for greater model quality. See project website [here](https://openvla-oft.github.io/). -- [2025-01-16] The FAST action tokenizer was recently released! Compared to vanilla OpenVLA-style 256-bin action discretization, FAST allows action chunks to be compressed into fewer tokens, speeding up inference by up to 15x when using discrete robot actions. See project website [here](https://www.physicalintelligence.company/research/fast). -- [2024-10-15] Added a [VLA Performance Troubleshooting](#vla-performance-troubleshooting) section to the README with best practices for debugging poor VLA performance after fine-tuning. -- [2024-09-04] Added LIBERO simulation benchmark fine-tuning experiments to paper (see v2 on [arXiv](https://arxiv.org/abs/2406.09246)); - added instructions for reproducing OpenVLA results in [LIBERO Simulation Benchmark Evaluations](#libero-simulation-benchmark-evaluations) section -- [2024-08-14] Added new section, [Evaluating OpenVLA](#evaluating-openvla), with instructions for running BridgeData V2 WidowX robot evals -- [2024-07-08] Added new sections: [Fine-Tuning OpenVLA via LoRA](#fine-tuning-openvla-via-lora), [Fully Fine-Tuning OpenVLA](#fully-fine-tuning-openvla) -- [2024-06-13] Initial release +## System Requirements -
+Inference: +* 1 GPU with ~16 GB VRAM for LIBERO sim benchmark tasks +* 1 GPU with ~18 GB VRAM for ALOHA robot tasks -A simple and scalable codebase for training and fine-tuning vision-language-action models (VLAs) for generalist robotic -manipulation: +Training: +* Between 1-8 GPUs with 27-80 GB, depending on the desired training setup (with default bfloat16 data type). See [this FAQ on our project website](https://openvla-oft.github.io/#train-compute) for details. -- **Different Dataset Mixtures**: We natively support arbitrary datasets in RLDS format, including arbitrary mixtures of - data from the [Open X-Embodiment Dataset](https://robotics-transformer-x.github.io/). -- **Easy Scaling**: Powered by PyTorch FSDP and Flash-Attention, we can quickly and efficiently train models from 1B - - 34B parameters, with easily adaptable model architectures. -- **Native Fine-Tuning Support**: Built-in support (with examples) for various forms of fine-tuning (full, - partial, LoRA). +## Quick Start -Built on top of [Prismatic VLMs](https://github.com/TRI-ML/prismatic-vlms). +First, set up a conda environment (see instructions in [SETUP.md](SETUP.md)). -## Getting Started - -To get started with loading and running OpenVLA models for inference, we provide a lightweight interface that leverages -HuggingFace `transformers` AutoClasses, with minimal dependencies. - -For example, to load `openvla-7b` for zero-shot instruction following in the -[BridgeData V2 environments](https://rail-berkeley.github.io/bridgedata/) with a WidowX robot: +Then, run the Python script below to download a pretrained OpenVLA-OFT checkpoint and run inference to generate an action chunk: ```python -# Install minimal dependencies (`torch`, `transformers`, `timm`, `tokenizers`, ...) -# > pip install -r https://raw.githubusercontent.com/openvla/openvla/main/requirements-min.txt -from transformers import AutoModelForVision2Seq, AutoProcessor -from PIL import Image - -import torch - -# Load Processor & VLA -processor = AutoProcessor.from_pretrained("openvla/openvla-7b", trust_remote_code=True) -vla = AutoModelForVision2Seq.from_pretrained( - "openvla/openvla-7b", - attn_implementation="flash_attention_2", # [Optional] Requires `flash_attn` - torch_dtype=torch.bfloat16, - low_cpu_mem_usage=True, - trust_remote_code=True -).to("cuda:0") - -# Grab image input & format prompt -image: Image.Image = get_from_camera(...) -prompt = "In: What action should the robot take to {}?\nOut:" - -# Predict Action (7-DoF; un-normalize for BridgeData V2) -inputs = processor(prompt, image).to("cuda:0", dtype=torch.bfloat16) -action = vla.predict_action(**inputs, unnorm_key="bridge_orig", do_sample=False) - -# Execute... -robot.act(action, ...) +import pickle +from experiments.robot.libero.run_libero_eval import GenerateConfig +from experiments.robot.openvla_utils import get_action_head, get_processor, get_proprio_projector, get_vla, get_vla_action +from prismatic.vla.constants import NUM_ACTIONS_CHUNK, PROPRIO_DIM + +# Instantiate config (see class GenerateConfig in experiments/robot/libero/run_libero_eval.py for definitions) +cfg = GenerateConfig( + pretrained_checkpoint = "moojink/openvla-7b-oft-finetuned-libero-spatial", + use_l1_regression = True, + use_diffusion = False, + use_film = False, + num_images_in_input = 2, + use_proprio = True, + load_in_8bit = False, + load_in_4bit = False, + center_crop = True, + num_open_loop_steps = NUM_ACTIONS_CHUNK, + unnorm_key = "libero_spatial_no_noops", +) + +# Load OpenVLA-OFT policy and inputs processor +vla = get_vla(cfg) +processor = get_processor(cfg) + +# Load MLP action head to generate continuous actions (via L1 regression) +action_head = get_action_head(cfg, llm_dim=vla.llm_dim) + +# Load proprio projector to map proprio to language embedding space +proprio_projector = get_proprio_projector(cfg, llm_dim=vla.llm_dim, proprio_dim=PROPRIO_DIM) + +# Load sample observation: +# observation (dict): { +# "full_image": primary third-person image, +# "wrist_image": wrist-mounted camera image, +# "state": robot proprioceptive state, +# "task_description": task description, +# } +with open("experiments/robot/libero/sample_libero_spatial_observation.pkl", "rb") as file: + observation = pickle.load(file) + +# Generate robot action chunk (sequence of future actions) +actions = get_vla_action(cfg, vla, processor, observation, observation["task_description"], action_head, proprio_projector) +print("Generated action chunk:") +for act in actions: + print(act) ``` -We also provide an [example script for fine-tuning OpenVLA models for new tasks and -embodiments](./vla-scripts/finetune.py); this script supports different fine-tuning modes -- including (quantized) -low-rank adaptation (LoRA) supported by [HuggingFace's PEFT library](https://huggingface.co/docs/peft/en/index). - -For deployment, we provide a lightweight script for [serving OpenVLA models over a REST API](./vla-scripts/deploy.py), -providing an easy way to integrate OpenVLA models into existing robot control stacks, -removing any requirement for powerful on-device compute. - -## Pretrained VLAs - -We release two OpenVLA models trained as part of our work, with checkpoints, configs, and model cards available [on our -HuggingFace page](https://huggingface.co/openvla): -- [`openvla-7b`](https://huggingface.co/openvla/openvla-7b): The flagship model from our paper, trained from - the Prismatic `prism-dinosiglip-224px` VLM (based on a fused DINOv2 and SigLIP vision backbone, and Llama-2 LLM). - Trained on a large mixture of datasets from Open X-Embodiment spanning 970K trajectories - ([mixture details - see "Open-X Magic Soup++"](./prismatic/vla/datasets/rlds/oxe/mixtures.py)). -- [`openvla-v01-7b`](https://huggingface.co/openvla/openvla-7b-v01): An early model used during development, trained from - the Prismatic `siglip-224px` VLM (singular SigLIP vision backbone, and a Vicuña v1.5 LLM). Trained on the same mixture - of datasets as [Octo](https://github.com/octo-models/octo), but for significantly fewer GPU hours than our final model - ([mixture details - see "Open-X Magic Soup"](./prismatic/vla/datasets/rlds/oxe/mixtures.py)). - -**Explicit Notes on Model Licensing & Commercial Use**: While all code in this repository is released under an MIT -License, our pretrained models may inherit restrictions from the underlying base models we use. Specifically, both the -above models are derived from Llama-2, and as such are subject to the -[Llama Community License](https://ai.meta.com/llama/license/). - ---- - ## Installation -> **Note**: These installation instructions are for full-scale pretraining (and distributed fine-tuning); if looking to - just run inference with OpenVLA models (or perform lightweight fine-tuning), see instructions above! - -This repository was built using Python 3.10, but should be backwards compatible with any Python >= 3.8. We require -PyTorch 2.2.* -- installation instructions [can be found here](https://pytorch.org/get-started/locally/). The latest -version of this repository was developed and thoroughly tested with: - - PyTorch 2.2.0, torchvision 0.17.0, transformers 4.40.1, tokenizers 0.19.1, timm 0.9.10, and flash-attn 2.5.5 - -**[5/21/24] Note**: Following reported regressions and breaking changes in later versions of `transformers`, `timm`, and -`tokenizers` we explicitly pin the above versions of the dependencies. We are working on implementing thorough tests, -and plan on relaxing these constraints as soon as we can. - -Use the setup commands below to get started: - -```bash -# Create and activate conda environment -conda create -n openvla python=3.10 -y -conda activate openvla - -# Install PyTorch. Below is a sample command to do this, but you should check the following link -# to find installation instructions that are specific to your compute platform: -# https://pytorch.org/get-started/locally/ -conda install pytorch torchvision torchaudio pytorch-cuda=12.4 -c pytorch -c nvidia -y # UPDATE ME! - -# Clone and install the openvla repo -git clone https://github.com/openvla/openvla.git -cd openvla -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 -``` - -If you run into any problems during the installation process, please file a GitHub Issue. - -**Note:** See `vla-scripts/` for full training and verification scripts for OpenVLA models. Note that `scripts/` is -mostly a holdover from the original (base) `prismatic-vlms` repository, with support for training and evaluating -visually-conditioned language models; while you can use this repo to train VLMs AND VLAs, note that trying to generate -language (via `scripts/generate.py`) with existing OpenVLA models will not work (as we only train current OpenVLA models -to generate actions, and actions alone). - -## Fine-Tuning OpenVLA via LoRA - -**(2025-03-03 Update: We recommend trying the new OFT recipe for fine-tuning OpenVLA to produce faster and more successful policies. See project website [here](https://openvla-oft.github.io/).)** - -In this section, we discuss fine-tuning OpenVLA using Low-Rank Adaptation (LoRA) via the Hugging Face `transformers` library, -which is recommended if you do not have sufficient compute to fully fine-tune a 7B-parameter model. The main script for LoRA -fine-tuning is `vla-scripts/finetune.py`. (If you instead wish to do full fine-tuning, please see the -[Fully Fine-Tuning OpenVLA](#fully-fine-tuning-openvla) section.) - -Below we show an example of how you can fine-tune the main OpenVLA checkpoint ([`openvla-7b`](https://huggingface.co/openvla/openvla-7b)) -via LoRA. Here we fine-tune on [BridgeData V2](https://rail-berkeley.github.io/bridgedata/) using a single A100 -GPU with 80 GB VRAM. (You can also fine-tune with a smaller GPU, as long as it has at least ~27 GB of memory, -by modifying the batch size.) - -First, download the BridgeData V2 dataset: - -```bash -# Change directory to your base datasets folder -cd - -# Download the full dataset (124 GB) -wget -r -nH --cut-dirs=4 --reject="index.html*" https://rail.eecs.berkeley.edu/datasets/bridge_release/data/tfds/bridge_dataset/ - -# Rename the dataset to `bridge_orig` (NOTE: Omitting this step may lead to runtime errors later) -mv bridge_dataset bridge_orig -``` - -Now, launch the LoRA fine-tuning script, as shown below. Note that `--batch_size==16` with `--grad_accumulation_steps==1` -requires ~72 GB GPU memory. If you have a smaller GPU, you should reduce `--batch_size` and increase `--grad_accumulation_steps` -to maintain an effective batch size that is large enough for stable training. If you have multiple GPUs and wish to train via -PyTorch Distributed Data Parallel (DDP), simply set `--nproc-per-node` in the `torchrun` command below to the number of available GPUs. - -```bash -torchrun --standalone --nnodes 1 --nproc-per-node 1 vla-scripts/finetune.py \ - --vla_path "openvla/openvla-7b" \ - --data_root_dir \ - --dataset_name bridge_orig \ - --run_root_dir \ - --adapter_tmp_dir \ - --lora_rank 32 \ - --batch_size 16 \ - --grad_accumulation_steps 1 \ - --learning_rate 5e-4 \ - --image_aug \ - --wandb_project \ - --wandb_entity \ - --save_steps -``` - -Note: If you set `--image_aug==False` in the command above, you will observe nearly 100% `action_accuracy` in the training logs, -since the [`openvla-7b`](https://huggingface.co/openvla/openvla-7b) model is already pretrained (without augmentations) on a -superset of datasets that includes BridgeData V2. - -To LoRA fine-tune on a different dataset, you can download the dataset from the [Open X-Embodiment (OXE)](https://robotics-transformer-x.github.io/) -mixture (see [this custom script](https://github.com/moojink/rlds_dataset_mod/blob/main/prepare_open_x.sh) for an example of how to download datasets -from OXE). Alternatively, if you have a custom dataset that is not part of OXE, you can either (a) convert the dataset to the RLDS format which is -compatible with our fine-tuning script (see [this repo](https://github.com/kpertsch/rlds_dataset_builder) for instructions on this), or (b) use your own -custom PyTorch Dataset wrapper (see comments in `vla-scripts/finetune.py` for instructions). We recommend option (a) for most users; the RLDS dataset and -dataloader are tested more extensively since we used these for all of our pretraining and fine-tuning experiments. - -For option (a), after you converted your dataset to RLDS, you need to register it with our data loader, by registering a dataset -config [here](prismatic/vla/datasets/rlds/oxe/configs.py#L54) and a dataset transform function [here](prismatic/vla/datasets/rlds/oxe/transforms.py#L828). - -Once you have integrated your new dataset, you can launch LoRA fine-tuning with the same `vla-scripts/finetune.py` script above. If you run into any issues, -please visit the [VLA Troubleshooting](#vla-troubleshooting) section or search for a similar issue in the [OpenVLA GitHub Issues page](https://github.com/openvla/openvla/issues?q=) -(including "Closed" issues). If you cannot find a similar issue there, feel free to create a new issue. - -## Fully Fine-Tuning OpenVLA - -**(2025-03-03 Update: We recommend trying the new OFT recipe for fine-tuning OpenVLA to produce faster and more successful policies. See project website [here](https://openvla-oft.github.io/).)** - -In this section, we discuss fully fine-tuning OpenVLA (all 7.5 billion parameters) via native PyTorch Fully Sharded Data Parallel (FSDP) -using the [Prismatic VLMs](https://github.com/TRI-ML/prismatic-vlms) training script. Full fine-tuning is more advanced/involved and is only recommended -if you have sufficient compute (e.g., a full node of 8 A100 GPUs) and if LoRA fine-tuning is insufficient for your use case (e.g., if the fine-tuning distribution -varies drastically from the pretraining distribution). Otherwise, we recommend that you try parameter-efficient fine-tuning via LoRA, which is described in the -[Fine-Tuning OpenVLA via LoRA](#fine-tuning-openvla-via-lora) section. - -For full fine-tuning, you will need to download [a different version of the OpenVLA model checkpoint](https://huggingface.co/openvla/openvla-7b-prismatic) that is compatible -with the Prismatic VLMs codebase, which we built on top of to develop the OpenVLA model. You can download this Prismatic-compatible OpenVLA checkpoint using the git commands below -(alternatively, you can download via the [Hugging Face CLI](https://huggingface.co/docs/huggingface_hub/main/en/guides/cli)): - -```bash -# Change directory to your base model checkpoints folder -cd - -# Download checkpoint (30 GB) -- may take a few minutes -git clone git@hf.co:openvla/openvla-7b-prismatic - -# If the command above did not download the full checkpoint, -# manually fetch it via git Large File Storage (LFS) -# Note: You may have to configure an SSH key for this to work -cd openvla-7b-prismatic -git lfs fetch --all -``` - -We show how you can fully fine-tune OpenVLA on [BridgeData V2](https://rail-berkeley.github.io/bridgedata/) using a single node with 8 GPUs. If you wish to -use a different number of GPUs (or nodes), you can modify the VLA training configuration in [`prismatic/conf/vla.py`](prismatic/conf/vla.py). - -Download the BridgeData V2 dataset: - -```bash -# Change directory to your base datasets folder -cd - -# Download the full dataset (124 GB) -wget -r -nH --cut-dirs=4 --reject="index.html*" https://rail.eecs.berkeley.edu/datasets/bridge_release/data/tfds/bridge_dataset/ - -# Rename the dataset to `bridge_orig` (NOTE: Omitting this step may lead to runtime errors later) -mv bridge_dataset bridge_orig -``` - -Next, create a [Hugging Face user access token](https://huggingface.co/docs/hub/en/security-tokens) and copy the token value (a string that starts with -`hf_...`) into a file named `.hf_token` at the root directory of this repo (`openvla/.hf_token`). - -```bash -# Go to openvla root directory -cd openvla - -# Copy HF token value into token file. Replace "hf_..." with your own token value! -# See: https://huggingface.co/docs/hub/en/security-tokens -echo hf_... >>> .hf_token -``` - -Now, launch the training script. If you wish to use a different number of nodes or GPUs, modify the VLA training configuration in -[`prismatic/conf/vla.py`](prismatic/conf/vla.py) and then change the `--nnodes` and `--nproc-per-node` arguments below accordingly. - -```bash -torchrun --standalone --nnodes 1 --nproc-per-node 8 vla-scripts/train.py \ - --pretrained_checkpoint \ - --vla.type prism-dinosiglip-224px+mx-bridge \ - --data_root_dir \ - --run_root_dir \ - --run_id \ - --image_aug \ - --wandb_project \ - --wandb_entity \ - --save_interval \ - --is_resume False -``` - -Note that the `--is_resume` argument is set to `False` above since we are fine-tuning a pretrained checkpoint rather than resuming a paused training run. - -If your training run gets paused and you wish to resume from the latest checkpoint, change `--pretrained_checkpoint` to the latest checkpoint path, -and then set `--is_resume==True` and specify `--resume_step` and `--resume_epoch` as the step and epoch number, respectively. For example, if you wish to -resume training from a checkpoint named `step-010000-epoch-20-loss=0.0160.pt`, you would set `is_resume==True`, `resume_step==10000`, and `resume_epoch==20`. - -Note: If you run the BridgeData V2 fine-tuning command above, you should observe nearly 100% Action Token Accuracy in the training logs, since the -[`openvla-7b`](https://huggingface.co/openvla/openvla-7b) model is already pretrained on a superset of datasets that includes BridgeData V2. - -To fully fine-tune OpenVLA on a different dataset, you can download the dataset from the [Open X-Embodiment (OXE)](https://robotics-transformer-x.github.io/) -mixture (see [this custom script](https://github.com/moojink/rlds_dataset_mod/blob/main/prepare_open_x.sh) for an example of how to download datasets from OXE). -Alternatively, if you have a custom dataset that is not part of OXE, you can convert the dataset to the RLDS format, which is compatible with our fine-tuning script -(see [this repo](https://github.com/kpertsch/rlds_dataset_builder) for instructions on this). After downloading/converting the dataset, you will need to modify the following files: - -* [`prismatic/conf/vla.py`](prismatic/conf/vla.py): Add a new training configuration by creating an experiment class, and then register it in the `VLARegistry` at the bottom of the file. - * Make sure to create a new unique `vla_id` for your fine-tuning run, and adjust some configuration variables as needed – e.g., `expected_world_size` (number of GPUs), - `per_device_batch_size` (batch size per GPU), `global_batch_size` (total batch size), `shuffle_buffer_size` (number of samples in shuffle buffer per GPU), etc. See comments - under the `VLAConfig` class at the top of the file to understand the purpose of each variable. -* [`prismatic/vla/datasets/rlds/oxe/mixtures.py`](prismatic/vla/datasets/rlds/oxe/mixtures.py): Define a new mixture for your fine-tuning mixture in the `OXE_NAMED_MIXTURES` dictionary. -* [`prismatic/vla/datasets/rlds/oxe/transforms.py`](prismatic/vla/datasets/rlds/oxe/transforms.py): Define a new dataset transform function for your fine-tuning dataset, and add it to the -`OXE_STANDARDIZATION_TRANSFORMS` registry at the bottom of the file. -* [`prismatic/vla/datasets/rlds/oxe/configs.py`](prismatic/vla/datasets/rlds/oxe/configs.py): Add a new configuration specifying your fine-tuning dataset's observation and action spaces -to the `OXE_DATASET_CONFIGS` dictionary. - -After completing the steps above, you can start full fine-tuning using the `vla-scripts/train.py` script. Make sure to set the `--vla.type` argument to the new `vla_id` that you added in `prismatic/conf/vla.py`. - -When you are finished with fine-tuning, you will need to convert the final model checkpoint to a version that is -compatible with the Hugging Face `transformers` library. See the [Converting Prismatic Models to Hugging Face](#converting-prismatic-models-to-hugging-face) section for instructions. - -If you run into any issues, please visit the [VLA Troubleshooting](#vla-troubleshooting) section or search for a similar issue in the -[OpenVLA GitHub Issues page](https://github.com/openvla/openvla/issues?q=) (including "Closed" issues). If you cannot find a similar issue there, feel free to create a new issue. - -### Converting Prismatic Models to Hugging Face - -If you have used the Prismatic VLMs codebase to train your model (e.g., if you did full fine-tuning of OpenVLA on a -new dataset), you will need to convert the final checkpoint to a version that is compatible with Hugging Face -`transformers` AutoClasses. We discuss how to do so in this section. - -Let's say your training run directory is `PRISMATIC_RUN_DIR` (e.g., `prism-dinosiglip-224px+mx-oxe-magic-soup-plus+n8+b32+x7`). -Inside this directory, there should be a directory called `checkpoints` which contains saved model checkpoints (e.g., -`step-295000-epoch-40-loss=0.2200.pt`). The Prismatic-to-Hugging-Face conversion script -([convert_openvla_weights_to_hf.py](vla-scripts/extern/convert_openvla_weights_to_hf.py)) expects a checkpoint file -named `latest-checkpoint.pt`. Therefore, you should first create a symbolic link called `latest-checkpoint.pt` that -points to the checkpoint file that you wish to convert: - -```bash -# Go to your Prismatic training run's `checkpoints` directory -cd PRISMATIC_RUN_DIR/checkpoints - -# Create symbolic link pointing to your checkpoint file -ln -s latest-checkpoint.pt -``` - -Then, launch the conversion script to convert the checkpoint from the Prismatic VLMs format to the Hugging Face format: - -```bash -python vla-scripts/extern/convert_openvla_weights_to_hf.py \ - --openvla_model_path_or_id \ - --output_hf_model_local_path -``` - -The command above will save the HF-compatible checkpoint in `output_hf_model_local_path`. Now you can load the checkpoint -with HF AutoClasses as normal, as shown below. Note that there is an additional necessary step to register the OpenVLA model -to HF AutoClasses before loading it because you are loading a locally saved checkpoint rather than one that is pushed to the -HF Hub (see [here](https://huggingface.co/docs/transformers/en/custom_models#registering-a-model-with-custom-code-to-the-auto-classes) -for details). - -```python -import torch -from transformers import AutoConfig, AutoImageProcessor, AutoModelForVision2Seq, AutoProcessor - -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 - -# Register OpenVLA model to HF AutoClasses (not needed if you pushed model to HF Hub) -AutoConfig.register("openvla", OpenVLAConfig) -AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) -AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) -AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) - -# Load Processor & VLA -processor = AutoProcessor.from_pretrained("", trust_remote_code=True) -vla = AutoModelForVision2Seq.from_pretrained( - "", - attn_implementation="flash_attention_2", # [Optional] Requires `flash_attn` - torch_dtype=torch.bfloat16, - low_cpu_mem_usage=True, - trust_remote_code=True, -).to("cuda:0") - -... -``` - -## Training VLAs from Scratch - -We provide full instructions and configurations for training VLA models on (arbitrary subsets of) the -[Open X-Embodiment (OXE) Dataset](https://robotics-transformer-x.github.io/). If you run in to any issues with -the following, see [VLA Troubleshooting](#vla-troubleshooting) below (or file a GitHub Issue). - -### VLA Pretraining Datasets - -We download and preprocess individual datasets from Open X-Embodiment in [RLDS format](https://github.com/google-research/rlds) following -[this custom script](https://github.com/moojink/rlds_dataset_mod/blob/main/prepare_open_x.sh). See -[mixtures.py](./prismatic/vla/datasets/rlds/oxe/mixtures.py) for the full list of component datasets (and mixture -weights) we use to train `openvla-7b`. -- **Important**: For the BridgeData V2 component, the version in OXE is out of date (as of 12/20/2023). Instead, - you should download the dataset from the [official website](https://rail.eecs.berkeley.edu/datasets/bridge_release/data/tfds/bridge_dataset/) and place it under the subdirectory `bridge_orig/`. - Replace any reference to `bridge` in the OXE code with `bridge_orig`. - -### VLA Configuration & Training Script - -The entry point for VLA training is [`vla-scripts/train.py`](vla-scripts/train.py). We use -[`draccus`](https://pypi.org/project/draccus) to provide a modular, dataclass-based interface for specifying VLA -training configurations; existing VLA configurations are in [`prismatic/conf/vla.py`](prismatic/conf/vla.py). You can -add your own training configuration and refer to it using the `--vla.type` command line argument. - -We use PyTorch Fully Sharded Data Parallel (FSDP) to distribute training across GPUs. Launch training via `torchrun`: - -```bash -# Train VLA on BridgeData V2 with the Prismatic DINO-SigLIP 224px Backbone on a Single Node (w/ 8 GPUs) -torchrun --standalone --nnodes 1 --nproc-per-node 8 vla-scripts/train.py \ - --vla.type "prism-dinosiglip-224px+mx-bridge" \ - --data_root_dir \ - --run_root_dir \ - --wandb_project "" \ - --wandb_entity "" -``` - -### VLA Troubleshooting - -The following are a list of known problems and corresponding fixes: - -```bash -FileNotFoundError: Failed to construct dataset "fractal20220817_data", builder_kwargs "{'data_dir': '/path/to/processed/datasets/'}": Could not load dataset info from fractal20220817_data/0.1.0/dataset_info.json -``` -- **Fix**: Downgrade `tensorflow-datasets` via `pip install tensorflow-datasets==4.9.3`. - - -```bash -AttributeError: 'DLataset' object has no attribute 'traj_map'. Did you mean: 'flat_map'? -``` -- **Fix**: Upgrade `dlimp` to the newest version. You may have to `--force-reinstall` like so: -`pip install --no-deps --force-reinstall git+https://github.com/moojink/dlimp_openvla` - ---- - -## Evaluating OpenVLA - -### BridgeData V2 WidowX Evaluations - -#### Setup - -Clone the [BridgeData V2 WidowX controller repo](https://github.com/rail-berkeley/bridge_data_robot) and install the `widowx_envs` package: - -```bash -git clone https://github.com/rail-berkeley/bridge_data_robot.git -cd bridge_data_robot -pip install -e widowx_envs -``` - -Additionally, install the [`edgeml`](https://github.com/youliangtan/edgeml) library: -```bash -git clone https://github.com/youliangtan/edgeml.git -cd edgeml -pip install -e . -``` - -Follow the instructions in the `bridge_data_robot` README to create the Bridge WidowX Docker container. - -#### Launching BridgeData V2 Evaluations - -There are multiple ways to run BridgeData V2 evaluations. We describe the server-client method below. - -In one Terminal window (e.g., in tmux), start the WidowX Docker container: - -```bash -cd bridge_data_robot -./generate_usb_config.sh -USB_CONNECTOR_CHART=$(pwd)/usb_connector_chart.yml docker compose up --build robonet -``` - -In a second Terminal window, run the WidowX robot server: - -```bash -cd bridge_data_robot -docker compose exec robonet bash -lic "widowx_env_service --server" -``` - -In a third Terminal window, run the OpenVLA policy evaluation script: - -```bash -cd openvla -python experiments/robot/bridge/run_bridgev2_eval.py \ - --model_family openvla \ - --pretrained_checkpoint openvla/openvla-7b -``` - -If you run into an error such as `ModuleNotFoundError: No module named 'moviepy.editor'`, you can work around it by fixing the -moviepy version to an older version, v1.0.3, in the bridge_data_robot repo's requirements.txt file -[here](https://github.com/rail-berkeley/bridge_data_robot/blob/main/widowx_envs/requirements.txt). I.e., simply replace `moviepy` -with `moviepy==1.0.3` in the requirements.txt file. Then, go back to the first step above and restart the WidowX Docker container; -it should be rebuilt with the older moviepy version. -``` - - -### LIBERO Simulation Benchmark Evaluations - -In the [updated OpenVLA paper (v2)](https://arxiv.org/abs/2406.09246), we discuss fine-tuning OpenVLA -on a simulated benchmark, [LIBERO](https://libero-project.github.io/main.html), in Appendix E. -Please see the paper for details, such as how we modify the provided demonstration datasets to -improve the overall performance of all methods. - -We copy the results to the section below and then discuss how to reproduce the results for OpenVLA. - -#### OpenVLA Fine-Tuning Results - -| Method | LIBERO-Spatial | LIBERO-Object | LIBERO-Goal | LIBERO-Long | Average | -|--------|----------------|---------------|-------------|-------------|---------| -| Diffusion Policy from scratch | 78.3 ± 1.1% | **92.5 ± 0.7%** | 68.3 ± 1.2% | 50.5 ± 1.3% | 72.4 ± 0.7% | -| Octo fine-tuned | 78.9 ± 1.0% | 85.7 ± 0.9% | **84.6 ± 0.9%** | 51.1 ± 1.3% | 75.1 ± 0.6% | -| OpenVLA fine-tuned (ours) | **84.7 ± 0.9%** | 88.4 ± 0.8% | 79.2 ± 1.0% | **53.7 ± 1.3%** | **76.5 ± 0.6%** | - -Each success rate is the average over 3 random seeds x 500 rollouts each (10 tasks x 50 rollouts per task). - -#### LIBERO Setup - -Clone and install the [LIBERO repo](https://github.com/Lifelong-Robot-Learning/LIBERO): - -```bash -git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git -cd LIBERO -pip install -e . -``` - -Additionally, install other required packages: -```bash -cd openvla -pip install -r experiments/robot/libero/libero_requirements.txt -``` - -(Optional) To download the modified versions of the LIBERO datasets that we used in our fine-tuning -experiments, run the command below. This will download the LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, -and LIBERO-10 datasets in RLDS data format (~10 GB total). You can use these to fine-tune OpenVLA or -train other methods. This step is optional since we provide pretrained OpenVLA checkpoints below. -(Also, you can find the script we used to generate the modified datasets in raw HDF5 format -[here](experiments/robot/libero/regenerate_libero_dataset.py) and the code we used to convert these -datasets to the RLDS format [here](https://github.com/moojink/rlds_dataset_builder).) -```bash -git clone git@hf.co:datasets/openvla/modified_libero_rlds -``` - -#### Launching LIBERO Evaluations - -We fine-tuned OpenVLA via LoRA (r=32) on four LIBERO task suites independently: LIBERO-Spatial, LIBERO-Object, LIBERO-Goal, and LIBERO-10 (also called LIBERO-Long). -The four checkpoints are available on Hugging Face: -* [openvla/openvla-7b-finetuned-libero-spatial](https://huggingface.co/openvla/openvla-7b-finetuned-libero-spatial) -* [openvla/openvla-7b-finetuned-libero-object](https://huggingface.co/openvla/openvla-7b-finetuned-libero-object) -* [openvla/openvla-7b-finetuned-libero-goal](https://huggingface.co/openvla/openvla-7b-finetuned-libero-goal) -* [openvla/openvla-7b-finetuned-libero-10](https://huggingface.co/openvla/openvla-7b-finetuned-libero-10) - -To start evaluation with one of these checkpoints, run one of the commands below. Each will automatically download the appropriate checkpoint listed above. - -```bash -# Launch LIBERO-Spatial evals -python experiments/robot/libero/run_libero_eval.py \ - --model_family openvla \ - --pretrained_checkpoint openvla/openvla-7b-finetuned-libero-spatial \ - --task_suite_name libero_spatial \ - --center_crop True - -# Launch LIBERO-Object evals -python experiments/robot/libero/run_libero_eval.py \ - --model_family openvla \ - --pretrained_checkpoint openvla/openvla-7b-finetuned-libero-object \ - --task_suite_name libero_object \ - --center_crop True - -# Launch LIBERO-Goal evals -python experiments/robot/libero/run_libero_eval.py \ - --model_family openvla \ - --pretrained_checkpoint openvla/openvla-7b-finetuned-libero-goal \ - --task_suite_name libero_goal \ - --center_crop True - -# Launch LIBERO-10 (LIBERO-Long) evals -python experiments/robot/libero/run_libero_eval.py \ - --model_family openvla \ - --pretrained_checkpoint openvla/openvla-7b-finetuned-libero-10 \ - --task_suite_name libero_10 \ - --center_crop True -``` - -Notes: -* The evaluation script will run 500 trials by default (10 tasks x 50 episodes each). You can modify the number of - trials per task by setting `--num_trials_per_task`. You can also change the random seed via `--seed`. -* **NOTE: Setting `--center_crop True` is important** because we fine-tuned OpenVLA with random crop augmentations - (we took a random crop with 90% area in every training sample, so at test time we simply take the center 90% crop). -* The evaluation script logs results locally. You can also log results in Weights & Biases - by setting `--use_wandb True` and specifying `--wandb_project ` and `--wandb_entity `. -* The results reported in our paper were obtained using **Python 3.10.13, PyTorch 2.2.0, transformers 4.40.1, and - flash-attn 2.5.5** on an **NVIDIA A100 GPU**, averaged over three random seeds. Please stick to these package versions. - Note that results may vary slightly if you use a different GPU for evaluation due to GPU nondeterminism in large models - (though we have tested that results were consistent across different machines with A100 GPUs). - -Please file a GitHub Issue if you run into any problems. - ---- - -## Repository Structure - -High-level overview of repository/project file-tree: - -+ `prismatic` - Package source; provides core utilities for model loading, training, data preprocessing, etc. -+ `vla-scripts/` - Core scripts for training, fine-tuning, and deploying VLAs. -+ `experiments/` - Code for evaluating OpenVLA policies in robot environments. -+ `LICENSE` - All code is made available under the MIT License; happy hacking! -+ `Makefile` - Top-level Makefile (by default, supports linting - checking & auto-fix); extend as needed. -+ `pyproject.toml` - Full project configuration details (including dependencies), as well as tool configurations. -+ `README.md` - You are here! - ---- - - -# VLA Performance Troubleshooting - -In this section we cover best practices for debugging poor VLA performance after fine-tuning on your target domain robot dataset. +See [SETUP.md](SETUP.md) for instructions on setting up the conda environment. -**Note**: OpenVLA typically requires fine-tuning on a small demonstration dataset (~100 demos) from your target domain robot. Out-of-the-box, it only works well on domains from the training dataset. +## Training and Evaluation -**Sanity checks**: -- replay the actions from a demonstration from your fine-tuning dataset and make sure that the robot can execute the task successfully (this ensures that your data collection pipeline is correct) -- once you fine-tuned a model, load the model in your inference pipeline (as if you would run it to control the robot), but feed images from the fine-tuning dataset into the model (pretending they come from the robot) and verify that you can reproduce the token accuracies / L1 errors from training (this ensures that your inference pipeline is correct) +See [LIBERO.md](LIBERO.md) for fine-tuning/evaluating on LIBERO simulation benchmark task suites. -**Best practices for fine-tuning data collection**: -If your setup passed the above two sanity checks, the issue may not be in model training, but in the data you fine-tuned the model with. Some best practices for data collection: -- *Collect at a control frequency around 5-10Hz.* OpenVLA is not trained with action chunking, empirically the model struggles with high-frequency data. If your robot setup uses a high-frequency controller (eg 50 Hz), consider downsampling your actions to 5Hz. Verify first that your robot can still solve the task when using 5Hz actions (ie repeat sanity check (1) above with 5Hz actions) -- *Avoid pauses / small actions during data collection.* Because OpenVLA is trained without action chunking, the model can be sensitive to idle actions in the fine-tuning data. If your data contains steps in which the robot barely moves, the model may "get stuck" in these steps at inference time. Try to collect fine-tuning demonstrations with continuous, slow movement. -- *Ensure sufficient data coverage.* If you plan to test the model with some variation, e.g. different initial positions of objects, make sure that your fine-tuning data contains sufficient diversity of such conditions as well, e.g. shows demonstrations with diverse initial conditions. -- *Use consistent task strategies during data collection.* This is not a hard constraint, but may make your life easier. Try to demonstrate tasks in consistent ways, e.g. approach objects from the same side, perform sub-steps in the same order even if they could be performed in arbitrary sequences. Being consistent gives you a less multi-modal fine-tuning dataset, which makes the modeling problem easier. +See [ALOHA.md](ALOHA.md) for fine-tuning/evaluating on real-world ALOHA robot tasks. +## Support ---- +If you run into any issues, please open a new GitHub issue. If you do not receive a response within 2 business days, please email Moo Jin Kim (moojink@cs.stanford.edu) to bring the issue to his attention. -#### Citation +## Citation -If you find our code or models useful in your work, please cite [our paper](https://arxiv.org/abs/2406.09246): +If you use our code in your work, please cite [our paper](https://arxiv.org/abs/2502.19645): ```bibtex -@article{kim24openvla, - title={OpenVLA: An Open-Source Vision-Language-Action Model}, - author={{Moo Jin} Kim and Karl Pertsch and Siddharth Karamcheti and Ted Xiao and Ashwin Balakrishna and Suraj Nair and Rafael Rafailov and Ethan Foster and Grace Lam and Pannag Sanketi and Quan Vuong and Thomas Kollar and Benjamin Burchfiel and Russ Tedrake and Dorsa Sadigh and Sergey Levine and Percy Liang and Chelsea Finn}, - journal = {arXiv preprint arXiv:2406.09246}, - year={2024} -} +@article{kim2025fine, + title={Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success}, + author={Kim, Moo Jin and Finn, Chelsea and Liang, Percy}, + journal={arXiv preprint arXiv:2502.19645}, + year={2025} +} ``` diff --git a/SETUP.md b/SETUP.md new file mode 100644 index 000000000..d4d7c72c7 --- /dev/null +++ b/SETUP.md @@ -0,0 +1,24 @@ +# Setup Instructions + +## Set Up Conda Environment + +```bash +# Create and activate conda environment +conda create -n openvla-oft python=3.10 -y +conda activate openvla-oft + +# 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 +``` \ No newline at end of file diff --git a/experiments/robot/aloha/aloha_utils.py b/experiments/robot/aloha/aloha_utils.py new file mode 100644 index 000000000..7acb0f24e --- /dev/null +++ b/experiments/robot/aloha/aloha_utils.py @@ -0,0 +1,85 @@ +"""Utils for evaluating policies in real-world ALOHA environments.""" + +import os + +import imageio +import numpy as np +from PIL import Image + +from experiments.robot.aloha.real_env import make_real_env +from experiments.robot.robot_utils import ( + DATE, + DATE_TIME, +) + + +def get_next_task_label(task_label): + """Prompt the user to input the next task.""" + if task_label == "": + user_input = "" + while user_input == "": + user_input = input("Enter the task name: ") + task_label = user_input + else: + user_input = input("Enter the task name (or leave blank to repeat the previous task): ") + if user_input == "": + pass # Do nothing -> Let task_label be the same + else: + task_label = user_input + print(f"Task: {task_label}") + return task_label + + +def get_aloha_env(): + """Initializes and returns the ALOHA environment.""" + env = make_real_env(init_node=True) + return env + + +def resize_image_for_preprocessing(img): + """ + Takes numpy array corresponding to a single image and resizes to 256x256, exactly as done + in the ALOHA data preprocessing script, which is used before converting the dataset to RLDS. + """ + ALOHA_PREPROCESS_SIZE = 256 + img = np.array( + Image.fromarray(img).resize((ALOHA_PREPROCESS_SIZE, ALOHA_PREPROCESS_SIZE), resample=Image.BICUBIC) + ) # BICUBIC is default; specify explicitly to make it clear + return img + + +def get_aloha_image(obs): + """Extracts third-person image from observations and preprocesses it.""" + # obs: dm_env._environment.TimeStep + img = obs.observation["images"]["cam_high"] + img = resize_image_for_preprocessing(img) + return img + + +def get_aloha_wrist_images(obs): + """Extracts both wrist camera images from observations and preprocesses them.""" + # obs: dm_env._environment.TimeStep + left_wrist_img = obs.observation["images"]["cam_left_wrist"] + right_wrist_img = obs.observation["images"]["cam_right_wrist"] + left_wrist_img = resize_image_for_preprocessing(left_wrist_img) + right_wrist_img = resize_image_for_preprocessing(right_wrist_img) + return left_wrist_img, right_wrist_img + + +def save_rollout_video(rollout_images, idx, success, task_description, log_file=None, notes=None): + """Saves an MP4 replay of an episode.""" + rollout_dir = f"./rollouts/{DATE}" + os.makedirs(rollout_dir, exist_ok=True) + processed_task_description = task_description.lower().replace(" ", "_").replace("\n", "_").replace(".", "_")[:50] + filetag = f"{rollout_dir}/{DATE_TIME}--openvla--episode={idx}--success={success}--task={processed_task_description}" + if notes is not None: + filetag += f"--{notes}" + mp4_path = f"{filetag}.mp4" + video_writer = imageio.get_writer(mp4_path, fps=25) + for img in rollout_images: + video_writer.append_data(img) + video_writer.close() + print(f"Saved rollout MP4 at path {mp4_path}") + if log_file is not None: + log_file.write(f"Saved rollout MP4 at path {mp4_path}\n") + return mp4_path diff --git a/experiments/robot/aloha/constants.py b/experiments/robot/aloha/constants.py new file mode 100644 index 000000000..5599e3590 --- /dev/null +++ b/experiments/robot/aloha/constants.py @@ -0,0 +1,100 @@ +### Task parameters + +DATA_DIR = '/scr2/moojink/data/aloha1/' +TASK_CONFIGS = { + # fold shorts + 'fold_shorts':{ + 'dataset_dir': DATA_DIR + '/fold_shorts', + 'num_episodes': 20, + 'episode_len': 1000, + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'] + }, + # fold shirt + 'fold_shirt':{ + 'dataset_dir': DATA_DIR + '/fold_shirt', + 'num_episodes': 30, + 'episode_len': 1250, + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'] + }, + # scoop X into bowl + 'scoop_raisins_into_bowl':{ + 'dataset_dir': DATA_DIR + '/scoop_raisins_into_bowl', + 'num_episodes': 15, + 'episode_len': 900, + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'] + }, + 'scoop_almonds_and_green_M&Ms_into_bowl':{ + 'dataset_dir': DATA_DIR + '/scoop_almonds_and_green_M&Ms_into_bowl', + 'num_episodes': 15, + 'episode_len': 900, + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'] + }, + 'scoop_pretzels_into_bowl':{ + 'dataset_dir': DATA_DIR + '/scoop_pretzels_into_bowl', + 'num_episodes': 15, + 'episode_len': 900, + 'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist'] + }, + # put X into pot + 'put_red_pepper_into_pot':{ + 'dataset_dir': DATA_DIR + '/put_red_pepper_into_pot', + 'num_episodes': 100, + 'episode_len': 400, + 'camera_names': ['cam_high', 'cam_low', 'cam_left_wrist', 'cam_right_wrist'] + }, + 'put_yellow_corn_into_pot':{ + 'dataset_dir': DATA_DIR + '/put_yellow_corn_into_pot', + 'num_episodes': 100, + 'episode_len': 400, + 'camera_names': ['cam_high', 'cam_low', 'cam_left_wrist', 'cam_right_wrist'] + }, + 'put_green_pepper_into_pot':{ + 'dataset_dir': DATA_DIR + '/put_green_pepper_into_pot', + 'num_episodes': 100, + 'episode_len': 400, + 'camera_names': ['cam_high', 'cam_low', 'cam_left_wrist', 'cam_right_wrist'] + }, +} + +### ALOHA fixed constants +DT = 0.04 # 1 / 0.04 -> 25 Hz +JOINT_NAMES = ["waist", "shoulder", "elbow", "forearm_roll", "wrist_angle", "wrist_rotate"] +START_ARM_POSE = [0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239, 0, -0.96, 1.16, 0, -0.3, 0, 0.02239, -0.02239] + +# Left finger position limits (qpos[7]), right_finger = -1 * left_finger +MASTER_GRIPPER_POSITION_OPEN = 0.02417 +MASTER_GRIPPER_POSITION_CLOSE = 0.01244 +PUPPET_GRIPPER_POSITION_OPEN = 0.05800 +PUPPET_GRIPPER_POSITION_CLOSE = 0.01844 + +# Gripper joint limits (qpos[6]) +MASTER_GRIPPER_JOINT_OPEN = 0.3083 # For ALOHA 1 +MASTER_GRIPPER_JOINT_CLOSE = -0.6842 # For ALOHA 1 +# MASTER_GRIPPER_JOINT_OPEN = -0.8 # For ALOHA 2 +# MASTER_GRIPPER_JOINT_CLOSE = -1.65 # For ALOHA 2 +PUPPET_GRIPPER_JOINT_OPEN = 1.4910 +PUPPET_GRIPPER_JOINT_CLOSE = -0.6213 + +############################ Helper functions ############################ + +MASTER_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_POSITION_CLOSE) / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) +PUPPET_GRIPPER_POSITION_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_POSITION_CLOSE) / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) +MASTER_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) + MASTER_GRIPPER_POSITION_CLOSE +PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + PUPPET_GRIPPER_POSITION_CLOSE +MASTER2PUPPET_POSITION_FN = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN(MASTER_GRIPPER_POSITION_NORMALIZE_FN(x)) + +MASTER_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) +PUPPET_GRIPPER_JOINT_NORMALIZE_FN = lambda x: (x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) +MASTER_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE +PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN = lambda x: x * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE +MASTER2PUPPET_JOINT_FN = lambda x: PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(MASTER_GRIPPER_JOINT_NORMALIZE_FN(x)) + +MASTER_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (MASTER_GRIPPER_POSITION_OPEN - MASTER_GRIPPER_POSITION_CLOSE) +PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN = lambda x: x / (PUPPET_GRIPPER_POSITION_OPEN - PUPPET_GRIPPER_POSITION_CLOSE) + +MASTER_POS2JOINT = lambda x: MASTER_GRIPPER_POSITION_NORMALIZE_FN(x) * (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE) + MASTER_GRIPPER_JOINT_CLOSE +MASTER_JOINT2POS = lambda x: MASTER_GRIPPER_POSITION_UNNORMALIZE_FN((x - MASTER_GRIPPER_JOINT_CLOSE) / (MASTER_GRIPPER_JOINT_OPEN - MASTER_GRIPPER_JOINT_CLOSE)) +PUPPET_POS2JOINT = lambda x: PUPPET_GRIPPER_POSITION_NORMALIZE_FN(x) * (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE) + PUPPET_GRIPPER_JOINT_CLOSE +PUPPET_JOINT2POS = lambda x: PUPPET_GRIPPER_POSITION_UNNORMALIZE_FN((x - PUPPET_GRIPPER_JOINT_CLOSE) / (PUPPET_GRIPPER_JOINT_OPEN - PUPPET_GRIPPER_JOINT_CLOSE)) + +MASTER_GRIPPER_JOINT_MID = (MASTER_GRIPPER_JOINT_OPEN + MASTER_GRIPPER_JOINT_CLOSE)/2 \ No newline at end of file diff --git a/experiments/robot/aloha/preprocess_split_aloha_data.py b/experiments/robot/aloha/preprocess_split_aloha_data.py new file mode 100644 index 000000000..8de07f232 --- /dev/null +++ b/experiments/robot/aloha/preprocess_split_aloha_data.py @@ -0,0 +1,260 @@ +""" +Preprocesses ALOHA dataset(s) and splits them into train/val sets. + +Preprocessing includes downsizing images from 480x640 to 256x256. +Splits happen at the episode level (not step level), which means that +an episode is treated as an atomic unit that entirely goes to either +the train set or val set. + +Original ALOHA data layout: + /PATH/TO/DATASET/dataset_name/ + - episode_0.hdf5 + - episode_1.hdf5 + - ... + - episode_N.hdf5 + +Preprocessed data layout (after running this script): + /PATH/TO/PREPROCESSED_DATASETS/dataset_name/ + - train/ + - episode_0.hdf5 + - episode_1.hdf5 + - ... + - episode_M.hdf5 + - val/ + - episode_0.hdf5 + - episode_1.hdf5 + - ... + - episode_K.hdf5 + + where N > M > K + +Example usage: + # "put X into pot" task + python experiments/robot/aloha/preprocess_split_aloha_data.py \ + --dataset_path /scr/moojink/data/aloha1_raw/put_green_pepper_into_pot/ \ + --out_base_dir /scr/moojink/data/aloha1_preprocessed/ \ + --percent_val 0.05 && \ + python experiments/robot/aloha/preprocess_split_aloha_data.py \ + --dataset_path /scr/moojink/data/aloha1_raw/put_red_pepper_into_pot/ \ + --out_base_dir /scr/moojink/data/aloha1_preprocessed/ \ + --percent_val 0.05 && \ + python experiments/robot/aloha/preprocess_split_aloha_data.py \ + --dataset_path /scr/moojink/data/aloha1_raw/put_yellow_corn_into_pot/ \ + --out_base_dir /scr/moojink/data/aloha1_preprocessed/ \ + --percent_val 0.05 +""" + +import argparse +import glob +import os +import random + +import h5py +import numpy as np +from PIL import Image +from tqdm import tqdm + + +def load_hdf5(demo_path): + """Loads single episode.""" + if not os.path.isfile(demo_path): + print(f"Dataset does not exist at \n{demo_path}\n") + exit() + + print(f"Loading {demo_path}...") + with h5py.File(demo_path, "r") as root: + is_sim = root.attrs["sim"] + qpos = root["/observations/qpos"][()] + qvel = root["/observations/qvel"][()] + effort = root["/observations/effort"][()] + action = root["/action"][()] + image_dict = dict() + for cam_name in root["/observations/images/"].keys(): + image_dict[cam_name] = root[f"/observations/images/{cam_name}"][()] + print(f"Loading episode complete: {demo_path}") + + return qpos, qvel, effort, action, image_dict, is_sim + + +def load_and_preprocess_all_episodes(demo_paths, out_dataset_dir): + """ + Loads and preprocesses all episodes. + Resizes all images in one episode before loading the next, to reduce memory usage. + """ + cam_names = ["cam_high", "cam_left_wrist", "cam_right_wrist"] + idx = 0 + for demo in tqdm(demo_paths): + qpos, qvel, effort, action, image_dict, is_sim = load_hdf5(demo) + # Save non-image info + episode_len = image_dict["cam_high"].shape[0] + # Resize all images + print("Resizing images in episode...") + for k in cam_names: + resized_images = [] + for i in range(episode_len): + resized_images.append( + np.array( + Image.fromarray(image_dict[k][i]).resize( + (args.img_resize_size, args.img_resize_size), resample=Image.BICUBIC + ) + ) # BICUBIC is default; specify explicitly to make it clear + ) + image_dict[k] = np.stack(resized_images) + print("Resizing images in episode complete!") + # Save preprocessed episode + data_dict = dict( + qpos=qpos, + qvel=qvel, + effort=effort, + action=action, + image_dict=image_dict, + is_sim=is_sim, + ) + save_new_hdf5(out_dataset_dir, data_dict, idx) + idx += 1 + + +def randomly_split(full_qpos, full_qvel, full_effort, full_action, full_image_dict, percent_val): + """Randomly splits dataset into train and validation sets.""" + # Create a list of episode indices + num_episodes_total = len(full_qpos) + indices = list(range(num_episodes_total)) + # Shuffle the episode indices + random.shuffle(indices) + # Create new lists using the shuffled indices + shuffled_qpos = [full_qpos[idx] for idx in indices] + shuffled_qvel = [full_qvel[idx] for idx in indices] + shuffled_effort = [full_effort[idx] for idx in indices] + shuffled_action = [full_action[idx] for idx in indices] + shuffled_image_dict = { + "cam_high": [], + "cam_left_wrist": [], + "cam_right_wrist": [], + } + for k in full_image_dict.keys(): + shuffled_image_dict[k] = [full_image_dict[k][idx] for idx in indices] + # Split into train and val sets + num_episodes_val = int(num_episodes_total * percent_val) + print(f"Total # steps: {num_episodes_total}; using {num_episodes_val} ({percent_val:.2f}%) for val set") + num_episodes_train = num_episodes_total - num_episodes_val + train_dict = dict( + qpos=shuffled_qpos[:num_episodes_train], + qvel=shuffled_qvel[:num_episodes_train], + effort=shuffled_effort[:num_episodes_train], + action=shuffled_action[:num_episodes_train], + image_dict=dict( + cam_high=shuffled_image_dict["cam_high"][:num_episodes_train], + cam_left_wrist=shuffled_image_dict["cam_left_wrist"][:num_episodes_train], + cam_right_wrist=shuffled_image_dict["cam_right_wrist"][:num_episodes_train], + ), + ) + val_dict = dict( + qpos=shuffled_qpos[num_episodes_train:], + qvel=shuffled_qvel[num_episodes_train:], + effort=shuffled_effort[num_episodes_train:], + action=shuffled_action[num_episodes_train:], + image_dict=dict( + cam_high=shuffled_image_dict["cam_high"][num_episodes_train:], + cam_left_wrist=shuffled_image_dict["cam_left_wrist"][num_episodes_train:], + cam_right_wrist=shuffled_image_dict["cam_right_wrist"][num_episodes_train:], + ), + ) + return train_dict, val_dict + + +def save_new_hdf5(out_dataset_dir, data_dict, episode_idx): + """Saves an HDF5 file for a new episode.""" + camera_names = data_dict["image_dict"].keys() + H, W, C = data_dict["image_dict"]["cam_high"][0].shape + out_path = os.path.join(out_dataset_dir, f"episode_{episode_idx}.hdf5") + # Save HDF5 with same structure as original demos (except that now we combine all episodes into one HDF5 file) + with h5py.File( + out_path, "w", rdcc_nbytes=1024**2 * 2 + ) as root: # Magic constant for rdcc_nbytes comes from ALOHA codebase + episode_len = data_dict["qpos"].shape[0] + root.attrs["sim"] = data_dict["is_sim"] + obs = root.create_group("observations") + _ = obs.create_dataset("qpos", (episode_len, 14)) + _ = obs.create_dataset("qvel", (episode_len, 14)) + _ = obs.create_dataset("effort", (episode_len, 14)) + root["/observations/qpos"][...] = data_dict["qpos"] + root["/observations/qvel"][...] = data_dict["qvel"] + root["/observations/effort"][...] = data_dict["effort"] + image = obs.create_group("images") + for cam_name in camera_names: + _ = image.create_dataset( + cam_name, + (episode_len, H, W, C), + dtype="uint8", + chunks=(1, H, W, C), + ) + root[f"/observations/images/{cam_name}"][...] = data_dict["image_dict"][cam_name] + _ = root.create_dataset("action", (episode_len, 14)) + root["/action"][...] = data_dict["action"] + # Compute and save *relative* actions as well + actions = data_dict["action"] + relative_actions = np.zeros_like(actions) + relative_actions[:-1] = actions[1:] - actions[:-1] # Relative actions are the changes in joint pos + relative_actions[-1] = relative_actions[-2] # Just copy the second-to-last action for the last action + _ = root.create_dataset("relative_action", (episode_len, 14)) + root["/relative_action"][...] = relative_actions + print(f"Saved dataset: {out_path}") + + +def main(args): + # Create directory to save preprocessed dataset (if it doesn't exist already) + os.makedirs(args.out_base_dir, exist_ok=True) + out_dataset_dir = os.path.join(args.out_base_dir, os.path.basename(args.dataset_path.rstrip("/"))) + os.makedirs(out_dataset_dir, exist_ok=True) + # Get list of filepaths of all episodes + all_demo_paths = glob.glob(os.path.join(args.dataset_path, "*.hdf5")) # List of HDF5 filepaths + all_demo_paths.sort() + # Create a list of episode indices + num_episodes_total = len(all_demo_paths) + indices = list(range(num_episodes_total)) + # Shuffle the episode indices + random.shuffle(indices) + # Split into train and val sets + num_episodes_val = int(num_episodes_total * args.percent_val) + print(f"Total # episodes: {num_episodes_total}; using {num_episodes_val} ({args.percent_val:.2f}%) for val set") + num_episodes_train = num_episodes_total - num_episodes_val + train_indices = indices[:num_episodes_train] + val_indices = indices[num_episodes_train:] + train_demo_paths = [all_demo_paths[i] for i in train_indices] + val_demo_paths = [all_demo_paths[i] for i in val_indices] + # Preprocess all episodes and save the result + out_dataset_dir_train = os.path.join(out_dataset_dir, "train") + out_dataset_dir_val = os.path.join(out_dataset_dir, "val") + os.makedirs(out_dataset_dir_train, exist_ok=True) + os.makedirs(out_dataset_dir_val, exist_ok=True) + load_and_preprocess_all_episodes(train_demo_paths, out_dataset_dir_train) + load_and_preprocess_all_episodes(val_demo_paths, out_dataset_dir_val) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--dataset_path", + required=True, + help="Path to raw ALOHA dataset directory. Example: /PATH/TO/USER/data/aloha_raw/put_green_pepper_into_pot/", + ) + parser.add_argument( + "--out_base_dir", + required=True, + help="Path to directory in which to save preprocessed dataset. Example: /PATH/TO/USER/data/aloha_preprocessed/", + ) + parser.add_argument( + "--percent_val", + type=float, + help="Percent of dataset to use as validation set (measured in episodes, not steps).", + default=0.05, + ) + parser.add_argument( + "--img_resize_size", + type=int, + help="Size to resize images to. Final images will be square (img_resize_size x img_resize_size pixels).", + default=256, + ) + args = parser.parse_args() + + main(args) diff --git a/experiments/robot/aloha/real_env.py b/experiments/robot/aloha/real_env.py new file mode 100644 index 000000000..f3f6c8f54 --- /dev/null +++ b/experiments/robot/aloha/real_env.py @@ -0,0 +1,213 @@ +import time +import numpy as np +import collections +import matplotlib.pyplot as plt +import dm_env + +from experiments.robot.aloha.constants import DT, START_ARM_POSE, MASTER_GRIPPER_JOINT_NORMALIZE_FN, PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN +from experiments.robot.aloha.constants import PUPPET_GRIPPER_POSITION_NORMALIZE_FN, PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN +from experiments.robot.aloha.constants import PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE +from experiments.robot.aloha.robot_utils import Recorder, ImageRecorder +from experiments.robot.aloha.robot_utils import setup_master_bot, setup_puppet_bot, move_arms, move_grippers +from interbotix_xs_modules.arm import InterbotixManipulatorXS +from interbotix_xs_msgs.msg import JointSingleCommand + +import IPython +e = IPython.embed + +class RealEnv: + """ + Environment for real robot bi-manual manipulation + Action space: [left_arm_qpos (6), # absolute joint position + left_gripper_positions (1), # normalized gripper position (0: close, 1: open) + right_arm_qpos (6), # absolute joint position + right_gripper_positions (1),] # normalized gripper position (0: close, 1: open) + + Observation space: {"qpos": Concat[ left_arm_qpos (6), # absolute joint position + left_gripper_position (1), # normalized gripper position (0: close, 1: open) + right_arm_qpos (6), # absolute joint position + right_gripper_qpos (1)] # normalized gripper position (0: close, 1: open) + "qvel": Concat[ left_arm_qvel (6), # absolute joint velocity (rad) + left_gripper_velocity (1), # normalized gripper velocity (pos: opening, neg: closing) + right_arm_qvel (6), # absolute joint velocity (rad) + right_gripper_qvel (1)] # normalized gripper velocity (pos: opening, neg: closing) + "images": {"cam_high": (480x640x3), # h, w, c, dtype='uint8' + "cam_low": (480x640x3), # h, w, c, dtype='uint8' + "cam_left_wrist": (480x640x3), # h, w, c, dtype='uint8' + "cam_right_wrist": (480x640x3)} # h, w, c, dtype='uint8' + """ + + def __init__(self, init_node, setup_robots=True): + self.puppet_bot_left = InterbotixManipulatorXS(robot_model="vx300s", group_name="arm", gripper_name="gripper", + robot_name=f'puppet_left', init_node=init_node) + self.puppet_bot_right = InterbotixManipulatorXS(robot_model="vx300s", group_name="arm", gripper_name="gripper", + robot_name=f'puppet_right', init_node=False) + if setup_robots: + self.setup_robots() + + self.recorder_left = Recorder('left', init_node=False) + self.recorder_right = Recorder('right', init_node=False) + self.image_recorder = ImageRecorder(init_node=False) + self.gripper_command = JointSingleCommand(name="gripper") + + def setup_robots(self): + setup_puppet_bot(self.puppet_bot_left) + setup_puppet_bot(self.puppet_bot_right) + + def get_qpos(self): + left_qpos_raw = self.recorder_left.qpos + right_qpos_raw = self.recorder_right.qpos + left_arm_qpos = left_qpos_raw[:6] + right_arm_qpos = right_qpos_raw[:6] + left_gripper_qpos = [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(left_qpos_raw[7])] # this is position not joint + right_gripper_qpos = [PUPPET_GRIPPER_POSITION_NORMALIZE_FN(right_qpos_raw[7])] # this is position not joint + return np.concatenate([left_arm_qpos, left_gripper_qpos, right_arm_qpos, right_gripper_qpos]) + + def get_qvel(self): + left_qvel_raw = self.recorder_left.qvel + right_qvel_raw = self.recorder_right.qvel + left_arm_qvel = left_qvel_raw[:6] + right_arm_qvel = right_qvel_raw[:6] + left_gripper_qvel = [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(left_qvel_raw[7])] + right_gripper_qvel = [PUPPET_GRIPPER_VELOCITY_NORMALIZE_FN(right_qvel_raw[7])] + return np.concatenate([left_arm_qvel, left_gripper_qvel, right_arm_qvel, right_gripper_qvel]) + + def get_effort(self): + left_effort_raw = self.recorder_left.effort + right_effort_raw = self.recorder_right.effort + left_robot_effort = left_effort_raw[:7] + right_robot_effort = right_effort_raw[:7] + return np.concatenate([left_robot_effort, right_robot_effort]) + + def get_images(self): + return self.image_recorder.get_images() + + def set_gripper_pose(self, left_gripper_desired_pos_normalized, right_gripper_desired_pos_normalized): + left_gripper_desired_joint = PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(left_gripper_desired_pos_normalized) + self.gripper_command.cmd = left_gripper_desired_joint + self.puppet_bot_left.gripper.core.pub_single.publish(self.gripper_command) + + right_gripper_desired_joint = PUPPET_GRIPPER_JOINT_UNNORMALIZE_FN(right_gripper_desired_pos_normalized) + self.gripper_command.cmd = right_gripper_desired_joint + self.puppet_bot_right.gripper.core.pub_single.publish(self.gripper_command) + + def _reset_joints(self): + reset_position = START_ARM_POSE[:6] + move_arms([self.puppet_bot_left, self.puppet_bot_right], [reset_position, reset_position], move_time=1) + + def _reset_gripper(self): + """Set to position mode and do position resets: first open then close. Then change back to PWM mode""" + move_grippers([self.puppet_bot_left, self.puppet_bot_right], [PUPPET_GRIPPER_JOINT_OPEN] * 2, move_time=0.5) + move_grippers([self.puppet_bot_left, self.puppet_bot_right], [PUPPET_GRIPPER_JOINT_CLOSE] * 2, move_time=1) + + def _get_obs(self): + obs = collections.OrderedDict() + obs['qpos'] = self.get_qpos() + obs['qvel'] = self.get_qvel() + obs['effort'] = self.get_effort() + obs['images'] = self.get_images() + return obs + + def get_observation(self, t=0): + step_type = dm_env.StepType.FIRST if t == 0 else dm_env.StepType.MID + return dm_env.TimeStep( + step_type=step_type, + reward=self.get_reward(), + discount=None, + observation=self._get_obs() + ) + + def get_reward(self): + return 0 + + def reset(self, fake=False): + if not fake: + # Reboot puppet robot gripper motors + self.puppet_bot_left.dxl.robot_reboot_motors("single", "gripper", True) + self.puppet_bot_right.dxl.robot_reboot_motors("single", "gripper", True) + self._reset_joints() + self._reset_gripper() + return dm_env.TimeStep( + step_type=dm_env.StepType.FIRST, + reward=self.get_reward(), + discount=None, + observation=self._get_obs()) + + def step(self, action): + state_len = int(len(action) / 2) + left_action = action[:state_len] + right_action = action[state_len:] + self.puppet_bot_left.arm.set_joint_positions(left_action[:6], blocking=False) + self.puppet_bot_right.arm.set_joint_positions(right_action[:6], blocking=False) + self.set_gripper_pose(left_action[-1], right_action[-1]) + time.sleep(DT) + return dm_env.TimeStep( + step_type=dm_env.StepType.MID, + reward=self.get_reward(), + discount=None, + observation=self._get_obs()) + + +def get_action(master_bot_left, master_bot_right): + action = np.zeros(14) # 6 joint + 1 gripper, for two arms + # Arm actions + action[:6] = master_bot_left.dxl.joint_states.position[:6] + action[7:7+6] = master_bot_right.dxl.joint_states.position[:6] + # Gripper actions + action[6] = MASTER_GRIPPER_JOINT_NORMALIZE_FN(master_bot_left.dxl.joint_states.position[6]) + action[7+6] = MASTER_GRIPPER_JOINT_NORMALIZE_FN(master_bot_right.dxl.joint_states.position[6]) + + return action + + +def make_real_env(init_node, setup_robots=True): + env = RealEnv(init_node, setup_robots) + return env + + +def test_real_teleop(): + """ + Test bimanual teleoperation and show image observations onscreen. + It first reads joint poses from both master arms. + Then use it as actions to step the environment. + The environment returns full observations including images. + + An alternative approach is to have separate scripts for teleoperation and observation recording. + This script will result in higher fidelity (obs, action) pairs + """ + + onscreen_render = True + render_cam = 'cam_left_wrist' + + # source of data + master_bot_left = InterbotixManipulatorXS(robot_model="wx250s", group_name="arm", gripper_name="gripper", + robot_name=f'master_left', init_node=True) + master_bot_right = InterbotixManipulatorXS(robot_model="wx250s", group_name="arm", gripper_name="gripper", + robot_name=f'master_right', init_node=False) + setup_master_bot(master_bot_left) + setup_master_bot(master_bot_right) + + # setup the environment + env = make_real_env(init_node=False) + ts = env.reset(fake=True) + episode = [ts] + # setup visualization + if onscreen_render: + ax = plt.subplot() + plt_img = ax.imshow(ts.observation['images'][render_cam]) + plt.ion() + + for t in range(1000): + action = get_action(master_bot_left, master_bot_right) + ts = env.step(action) + episode.append(ts) + + if onscreen_render: + plt_img.set_data(ts.observation['images'][render_cam]) + plt.pause(DT) + else: + time.sleep(DT) + + +if __name__ == '__main__': + test_real_teleop() diff --git a/experiments/robot/aloha/requirements_aloha.txt b/experiments/robot/aloha/requirements_aloha.txt new file mode 100644 index 000000000..c84c6d08c --- /dev/null +++ b/experiments/robot/aloha/requirements_aloha.txt @@ -0,0 +1,26 @@ +numpy<2 +draccus +torchvision +torch +pyquaternion +pyyaml +rospkg +pexpect +mujoco==2.3.7 +dm_control==1.0.14 +opencv-python +matplotlib +einops +packaging +h5py +traitlets +ipdb +IPython +modern_robotics +Pillow +termcolor +imageio[ffmpeg] +uvicorn +fastapi +requests +json_numpy diff --git a/experiments/robot/aloha/robot_utils.py b/experiments/robot/aloha/robot_utils.py new file mode 100644 index 000000000..82f080cc9 --- /dev/null +++ b/experiments/robot/aloha/robot_utils.py @@ -0,0 +1,187 @@ +import numpy as np +import time +from experiments.robot.aloha.constants import DT +from interbotix_xs_msgs.msg import JointSingleCommand + +import IPython +e = IPython.embed + +class ImageRecorder: + def __init__(self, init_node=True, is_debug=False): + from collections import deque + import rospy + from cv_bridge import CvBridge + from sensor_msgs.msg import Image + self.is_debug = is_debug + self.bridge = CvBridge() + self.camera_names = ['cam_high', 'cam_low', 'cam_left_wrist', 'cam_right_wrist'] + if init_node: + rospy.init_node('image_recorder', anonymous=True) + for cam_name in self.camera_names: + setattr(self, f'{cam_name}_image', None) + setattr(self, f'{cam_name}_secs', None) + setattr(self, f'{cam_name}_nsecs', None) + if cam_name == 'cam_high': + callback_func = self.image_cb_cam_high + elif cam_name == 'cam_low': + callback_func = self.image_cb_cam_low + elif cam_name == 'cam_left_wrist': + callback_func = self.image_cb_cam_left_wrist + elif cam_name == 'cam_right_wrist': + callback_func = self.image_cb_cam_right_wrist + else: + raise NotImplementedError + rospy.Subscriber(f"/usb_{cam_name}/image_raw", Image, callback_func) + if self.is_debug: + setattr(self, f'{cam_name}_timestamps', deque(maxlen=50)) + time.sleep(0.5) + + def image_cb(self, cam_name, data): + setattr(self, f'{cam_name}_image', self.bridge.imgmsg_to_cv2(data, desired_encoding='passthrough')) + setattr(self, f'{cam_name}_secs', data.header.stamp.secs) + setattr(self, f'{cam_name}_nsecs', data.header.stamp.nsecs) + # cv2.imwrite('/home/tonyzhao/Desktop/sample.jpg', cv_image) + if self.is_debug: + getattr(self, f'{cam_name}_timestamps').append(data.header.stamp.secs + data.header.stamp.secs * 1e-9) + + def image_cb_cam_high(self, data): + cam_name = 'cam_high' + return self.image_cb(cam_name, data) + + def image_cb_cam_low(self, data): + cam_name = 'cam_low' + return self.image_cb(cam_name, data) + + def image_cb_cam_left_wrist(self, data): + cam_name = 'cam_left_wrist' + return self.image_cb(cam_name, data) + + def image_cb_cam_right_wrist(self, data): + cam_name = 'cam_right_wrist' + return self.image_cb(cam_name, data) + + def get_images(self): + image_dict = dict() + for cam_name in self.camera_names: + image_dict[cam_name] = getattr(self, f'{cam_name}_image') + return image_dict + + def print_diagnostics(self): + def dt_helper(l): + l = np.array(l) + diff = l[1:] - l[:-1] + return np.mean(diff) + for cam_name in self.camera_names: + image_freq = 1 / dt_helper(getattr(self, f'{cam_name}_timestamps')) + print(f'{cam_name} {image_freq=:.2f}') + print() + +class Recorder: + def __init__(self, side, init_node=True, is_debug=False): + from collections import deque + import rospy + from sensor_msgs.msg import JointState + from interbotix_xs_msgs.msg import JointGroupCommand, JointSingleCommand + + self.secs = None + self.nsecs = None + self.qpos = None + self.effort = None + self.arm_command = None + self.gripper_command = None + self.is_debug = is_debug + + if init_node: + rospy.init_node('recorder', anonymous=True) + rospy.Subscriber(f"/puppet_{side}/joint_states", JointState, self.puppet_state_cb) + rospy.Subscriber(f"/puppet_{side}/commands/joint_group", JointGroupCommand, self.puppet_arm_commands_cb) + rospy.Subscriber(f"/puppet_{side}/commands/joint_single", JointSingleCommand, self.puppet_gripper_commands_cb) + if self.is_debug: + self.joint_timestamps = deque(maxlen=50) + self.arm_command_timestamps = deque(maxlen=50) + self.gripper_command_timestamps = deque(maxlen=50) + time.sleep(0.1) + + def puppet_state_cb(self, data): + self.qpos = data.position + self.qvel = data.velocity + self.effort = data.effort + self.data = data + if self.is_debug: + self.joint_timestamps.append(time.time()) + + def puppet_arm_commands_cb(self, data): + self.arm_command = data.cmd + if self.is_debug: + self.arm_command_timestamps.append(time.time()) + + def puppet_gripper_commands_cb(self, data): + self.gripper_command = data.cmd + if self.is_debug: + self.gripper_command_timestamps.append(time.time()) + + def print_diagnostics(self): + def dt_helper(l): + l = np.array(l) + diff = l[1:] - l[:-1] + return np.mean(diff) + + joint_freq = 1 / dt_helper(self.joint_timestamps) + arm_command_freq = 1 / dt_helper(self.arm_command_timestamps) + gripper_command_freq = 1 / dt_helper(self.gripper_command_timestamps) + + print(f'{joint_freq=:.2f}\n{arm_command_freq=:.2f}\n{gripper_command_freq=:.2f}\n') + +def get_arm_joint_positions(bot): + return bot.arm.core.joint_states.position[:6] + +def get_arm_gripper_positions(bot): + joint_position = bot.gripper.core.joint_states.position[6] + return joint_position + +def move_arms(bot_list, target_pose_list, move_time=1): + num_steps = int(move_time / DT) + curr_pose_list = [get_arm_joint_positions(bot) for bot in bot_list] + traj_list = [np.linspace(curr_pose, target_pose, num_steps) for curr_pose, target_pose in zip(curr_pose_list, target_pose_list)] + for t in range(num_steps): + for bot_id, bot in enumerate(bot_list): + bot.arm.set_joint_positions(traj_list[bot_id][t], blocking=False) + time.sleep(DT) + +def move_grippers(bot_list, target_pose_list, move_time): + gripper_command = JointSingleCommand(name="gripper") + num_steps = int(move_time / DT) + curr_pose_list = [get_arm_gripper_positions(bot) for bot in bot_list] + traj_list = [np.linspace(curr_pose, target_pose, num_steps) for curr_pose, target_pose in zip(curr_pose_list, target_pose_list)] + for t in range(num_steps): + for bot_id, bot in enumerate(bot_list): + gripper_command.cmd = traj_list[bot_id][t] + bot.gripper.core.pub_single.publish(gripper_command) + time.sleep(DT) + +def setup_puppet_bot(bot): + bot.dxl.robot_reboot_motors("single", "gripper", True) + bot.dxl.robot_set_operating_modes("group", "arm", "position") + bot.dxl.robot_set_operating_modes("single", "gripper", "current_based_position") + torque_on(bot) + +def setup_master_bot(bot): + bot.dxl.robot_set_operating_modes("group", "arm", "pwm") + bot.dxl.robot_set_operating_modes("single", "gripper", "current_based_position") + torque_off(bot) + +def set_standard_pid_gains(bot): + bot.dxl.robot_set_motor_registers("group", "arm", 'Position_P_Gain', 800) + bot.dxl.robot_set_motor_registers("group", "arm", 'Position_I_Gain', 0) + +def set_low_pid_gains(bot): + bot.dxl.robot_set_motor_registers("group", "arm", 'Position_P_Gain', 100) + bot.dxl.robot_set_motor_registers("group", "arm", 'Position_I_Gain', 0) + +def torque_off(bot): + bot.dxl.robot_torque_enable("group", "arm", False) + bot.dxl.robot_torque_enable("single", "gripper", False) + +def torque_on(bot): + bot.dxl.robot_torque_enable("group", "arm", True) + bot.dxl.robot_torque_enable("single", "gripper", True) \ No newline at end of file diff --git a/experiments/robot/aloha/run_aloha_eval.py b/experiments/robot/aloha/run_aloha_eval.py new file mode 100644 index 000000000..520f5af9a --- /dev/null +++ b/experiments/robot/aloha/run_aloha_eval.py @@ -0,0 +1,385 @@ +""" +run_aloha_eval.py + +Evaluates a model in a real-world ALOHA environment. +""" + +import logging +import os +import socket +import sys +import time +from collections import deque +from dataclasses import dataclass +from pathlib import Path +from typing import Optional, Union + +import draccus +import tqdm + +# Append current directory so that interpreter can find experiments.robot +sys.path.append(".") +from experiments.robot.aloha.aloha_utils import ( + get_aloha_env, + get_aloha_image, + get_aloha_wrist_images, + get_next_task_label, + save_rollout_video, +) +from experiments.robot.openvla_utils import ( + get_action_from_server, + resize_image_for_policy, +) +from experiments.robot.robot_utils import ( + DATE_TIME, + get_image_resize_size, + set_seed_everywhere, +) + +# Set up logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + handlers=[logging.StreamHandler()], +) +logger = logging.getLogger(__name__) + + +@dataclass +class GenerateConfig: + # fmt: off + + ################################################################################################################# + # Model-specific parameters + ################################################################################################################# + model_family: str = "openvla" # Model family + + center_crop: bool = True # Center crop? (if trained w/ random crop image aug) + num_open_loop_steps: int = 25 # Number of actions to execute open-loop before requerying policy + + use_vla_server: bool = True # Whether to query remote VLA server for actions + vla_server_url: Union[str, Path] = "" # Remote VLA server URL (set to 127.0.0.1 if on same machine) + + ################################################################################################################# + # ALOHA environment-specific parameters + ################################################################################################################# + num_rollouts_planned: int = 50 # Number of test rollouts + max_steps: int = 1500 # Max number of steps per rollout + use_relative_actions: bool = False # Whether to use relative actions (delta joint angles) + + ################################################################################################################# + # Utils + ################################################################################################################# + run_id_note: Optional[str] = None # Extra note to add to end of run ID for logging + local_log_dir: str = "./experiments/logs" # Local directory for eval logs + + seed: int = 7 # Random Seed (for reproducibility) + + # fmt: on + + +def validate_config(cfg: GenerateConfig) -> None: + """Validate configuration parameters.""" + assert cfg.use_vla_server, ( + "Must use VLA server (server-client interface) to query model and get actions! Please set --use_vla_server=True" + ) + + +def setup_logging(cfg: GenerateConfig): + """Set up logging to file.""" + # Create run ID + run_id = f"EVAL-{cfg.model_family}-{DATE_TIME}" + if cfg.run_id_note is not None: + run_id += f"--{cfg.run_id_note}" + + # Set up local logging + os.makedirs(cfg.local_log_dir, exist_ok=True) + local_log_filepath = os.path.join(cfg.local_log_dir, run_id + ".txt") + log_file = open(local_log_filepath, "w") + logger.info(f"Logging to local log file: {local_log_filepath}") + + return log_file, local_log_filepath, run_id + + +def log_message(message: str, log_file=None): + """Log a message to console and optionally to a log file.""" + print(message) + logger.info(message) + if log_file: + log_file.write(message + "\n") + log_file.flush() + + +def get_server_endpoint(cfg: GenerateConfig): + """Get the server endpoint for remote inference.""" + ip_address = socket.gethostbyname(cfg.vla_server_url) + return f"http://{ip_address}:8777/act" + + +def prepare_observation(obs, resize_size): + """Prepare observation for policy input.""" + # Get preprocessed images + img = get_aloha_image(obs) + left_wrist_img, right_wrist_img = get_aloha_wrist_images(obs) + + # Resize images to size expected by model + img_resized = resize_image_for_policy(img, resize_size) + left_wrist_img_resized = resize_image_for_policy(left_wrist_img, resize_size) + right_wrist_img_resized = resize_image_for_policy(right_wrist_img, resize_size) + + # Prepare observations dict + observation = { + "full_image": img_resized, + "left_wrist_image": left_wrist_img_resized, + "right_wrist_image": right_wrist_img_resized, + "state": obs.observation["qpos"], + } + + return observation, img_resized, left_wrist_img_resized, right_wrist_img_resized + + +def run_episode( + cfg: GenerateConfig, + env, + task_description: str, + server_endpoint: str, + resize_size, + log_file=None, +): + """Run a single episode in the ALOHA environment.""" + # Define control frequency + STEP_DURATION_IN_SEC = 1.0 / 25.0 + + # Reset environment + obs = env.reset() + + # Initialize action queue + action_queue = deque(maxlen=cfg.num_open_loop_steps) + + # Setup + t = 0 + curr_state = None + replay_images = [] + replay_images_resized = [] + replay_images_left_wrist_resized = [] + replay_images_right_wrist_resized = [] + + log_message("Prepare the scene, and then press Enter to begin...", log_file) + input() + + # Reset environment again to fetch first timestep observation + obs = env.reset() + + # Fetch initial robot state (but sleep first so that robot stops moving) + time.sleep(2) + curr_state = env.get_qpos() + + episode_start_time = time.time() + total_model_query_time = 0.0 + + try: + while t < cfg.max_steps: + # Get step start time (used to compute how much to sleep between steps) + step_start_time = time.time() + + # Get observation + obs = env.get_observation(t=t) + + # Save raw high camera image for replay video + replay_images.append(obs.observation["images"]["cam_high"]) + + # If action queue is empty, requery model + if len(action_queue) == 0: + # Prepare observation + observation, img_resized, left_wrist_resized, right_wrist_resized = prepare_observation(obs, resize_size) + observation["instruction"] = task_description + + # Save processed images for replay + replay_images_resized.append(img_resized) + replay_images_left_wrist_resized.append(left_wrist_resized) + replay_images_right_wrist_resized.append(right_wrist_resized) + + # Query model to get action + log_message("Requerying model...", log_file) + model_query_start_time = time.time() + actions = get_action_from_server(observation, server_endpoint) + actions = actions[: cfg.num_open_loop_steps] + total_model_query_time += time.time() - model_query_start_time + action_queue.extend(actions) + + # Get action from queue + action = action_queue.popleft() + log_message("-----------------------------------------------------", log_file) + log_message(f"t: {t}", log_file) + log_message(f"action: {action}", log_file) + + # Execute action in environment + if cfg.use_relative_actions: + # Get absolute joint angles from relative action + rel_action = action + target_state = curr_state + rel_action + obs = env.step(target_state.tolist()) + # Update current state (assume it is the commanded target state) + curr_state = target_state + else: + obs = env.step(action.tolist()) + t += 1 + + # Sleep until next timestep + step_elapsed_time = time.time() - step_start_time + if step_elapsed_time < STEP_DURATION_IN_SEC: + time_to_sleep = STEP_DURATION_IN_SEC - step_elapsed_time + log_message(f"Sleeping {time_to_sleep} sec...", log_file) + time.sleep(time_to_sleep) + + except (KeyboardInterrupt, Exception) as e: + if isinstance(e, KeyboardInterrupt): + log_message("\nCaught KeyboardInterrupt: Terminating episode early.", log_file) + else: + log_message(f"\nCaught exception: {e}", log_file) + + episode_end_time = time.time() + + # Get success feedback from user + user_input = input("Success? Enter 'y' or 'n': ") + success = True if user_input.lower() == "y" else False + + # Calculate episode statistics + episode_stats = { + "success": success, + "total_steps": t, + "model_query_time": total_model_query_time, + "episode_duration": episode_end_time - episode_start_time, + } + + return ( + episode_stats, + replay_images, + replay_images_resized, + replay_images_left_wrist_resized, + replay_images_right_wrist_resized, + ) + + +def save_episode_videos( + replay_images, + replay_images_resized, + replay_images_left_wrist, + replay_images_right_wrist, + episode_idx, + success, + task_description, + log_file=None, +): + """Save videos of the episode from different camera angles.""" + # Save main replay video + save_rollout_video(replay_images, episode_idx, success=success, task_description=task_description, log_file=log_file) + + # Save processed view videos + save_rollout_video( + replay_images_resized, + episode_idx, + success=success, + task_description=task_description, + log_file=log_file, + notes="resized", + ) + save_rollout_video( + replay_images_left_wrist, + episode_idx, + success=success, + task_description=task_description, + log_file=log_file, + notes="left_wrist_resized", + ) + save_rollout_video( + replay_images_right_wrist, + episode_idx, + success=success, + task_description=task_description, + log_file=log_file, + notes="right_wrist_resized", + ) + + +@draccus.wrap() +def eval_aloha(cfg: GenerateConfig) -> None: + """Main function to evaluate a trained policy in a real-world ALOHA environment.""" + # Validate configuration + validate_config(cfg) + + # Set random seed + set_seed_everywhere(cfg.seed) + + # Setup logging + log_file, local_log_filepath, run_id = setup_logging(cfg) + + # Get expected image dimensions + resize_size = get_image_resize_size(cfg) + + # Get ALOHA environment + env = get_aloha_env() + + # Get server endpoint for remote inference + server_endpoint = get_server_endpoint(cfg) + + # Initialize task description + task_description = "" + + # Start evaluation + num_rollouts_completed, total_successes = 0, 0 + + for episode_idx in tqdm.tqdm(range(cfg.num_rollouts_planned)): + # Get task description from user + task_description = get_next_task_label(task_description) + log_message(f"\nTask: {task_description}", log_file) + + log_message(f"Starting episode {num_rollouts_completed + 1}...", log_file) + + # Run episode + episode_stats, replay_images, replay_images_resized, replay_images_left_wrist, replay_images_right_wrist = ( + run_episode(cfg, env, task_description, server_endpoint, resize_size, log_file) + ) + + # Update counters + num_rollouts_completed += 1 + if episode_stats["success"]: + total_successes += 1 + + # Save videos + save_episode_videos( + replay_images, + replay_images_resized, + replay_images_left_wrist, + replay_images_right_wrist, + num_rollouts_completed, + episode_stats["success"], + task_description, + log_file, + ) + + # Log results + log_message(f"Success: {episode_stats['success']}", log_file) + log_message(f"# episodes completed so far: {num_rollouts_completed}", log_file) + log_message(f"# successes: {total_successes} ({total_successes / num_rollouts_completed * 100:.1f}%)", log_file) + log_message(f"Total model query time: {episode_stats['model_query_time']:.2f} sec", log_file) + log_message(f"Total episode elapsed time: {episode_stats['episode_duration']:.2f} sec", log_file) + + # Calculate final success rate + final_success_rate = float(total_successes) / float(num_rollouts_completed) if num_rollouts_completed > 0 else 0 + + # Log final results + log_message("\nFinal results:", log_file) + log_message(f"Total episodes: {num_rollouts_completed}", log_file) + log_message(f"Total successes: {total_successes}", log_file) + log_message(f"Overall success rate: {final_success_rate:.4f} ({final_success_rate * 100:.1f}%)", log_file) + + # Close log file + if log_file: + log_file.close() + + return final_success_rate + + +if __name__ == "__main__": + eval_aloha() diff --git a/experiments/robot/libero/libero_utils.py b/experiments/robot/libero/libero_utils.py index 70a5d7074..9d1f3fa1a 100644 --- a/experiments/robot/libero/libero_utils.py +++ b/experiments/robot/libero/libero_utils.py @@ -30,31 +30,17 @@ def get_libero_dummy_action(model_family: str): return [0, 0, 0, 0, 0, 0, -1] -def resize_image(img, resize_size): - """ - Takes numpy array corresponding to a single image and returns resized image as numpy array. - - NOTE (Moo Jin): To make input images in distribution with respect to the inputs seen at training time, we follow - the same resizing scheme used in the Octo dataloader, which OpenVLA uses for training. - """ - assert isinstance(resize_size, tuple) - # Resize to image size expected by model - img = tf.image.encode_jpeg(img) # Encode as JPEG, as done in RLDS dataset builder - img = tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8) # Immediately decode back - img = tf.image.resize(img, resize_size, method="lanczos3", antialias=True) - img = tf.cast(tf.clip_by_value(tf.round(img), 0, 255), tf.uint8) - img = img.numpy() +def get_libero_image(obs): + """Extracts third-person image from observations and preprocesses it.""" + img = obs["agentview_image"] + img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing return img -def get_libero_image(obs, resize_size): - """Extracts image from observations and preprocesses it.""" - assert isinstance(resize_size, int) or isinstance(resize_size, tuple) - if isinstance(resize_size, int): - resize_size = (resize_size, resize_size) - img = obs["agentview_image"] +def get_libero_wrist_image(obs): + """Extracts wrist camera image from observations and preprocesses it.""" + img = obs["robot0_eye_in_hand_image"] img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing - img = resize_image(img, resize_size) return img diff --git a/experiments/robot/libero/run_libero_eval.py b/experiments/robot/libero/run_libero_eval.py index 5c3f58178..ed370adaa 100644 --- a/experiments/robot/libero/run_libero_eval.py +++ b/experiments/robot/libero/run_libero_eval.py @@ -1,25 +1,16 @@ """ run_libero_eval.py -Runs a model in a LIBERO simulation environment. - -Usage: - # OpenVLA: - # IMPORTANT: Set `center_crop=True` if model is fine-tuned with augmentations - python experiments/robot/libero/run_libero_eval.py \ - --model_family openvla \ - --pretrained_checkpoint \ - --task_suite_name [ libero_spatial | libero_object | libero_goal | libero_10 | libero_90 ] \ - --center_crop [ True | False ] \ - --run_id_note \ - --use_wandb [ True | False ] \ - --wandb_project \ - --wandb_entity +Evaluates a trained policy in a LIBERO simulation benchmark task suite. """ +import json +import logging import os import sys +from collections import deque from dataclasses import dataclass +from enum import Enum from pathlib import Path from typing import Optional, Union @@ -36,10 +27,17 @@ get_libero_dummy_action, get_libero_env, get_libero_image, + get_libero_wrist_image, quat2axisangle, save_rollout_video, ) -from experiments.robot.openvla_utils import get_processor +from experiments.robot.openvla_utils import ( + get_action_head, + get_noisy_action_projector, + get_processor, + get_proprio_projector, + resize_image_for_policy, +) from experiments.robot.robot_utils import ( DATE_TIME, get_action, @@ -49,6 +47,35 @@ normalize_gripper_action, set_seed_everywhere, ) +from prismatic.vla.constants import NUM_ACTIONS_CHUNK + + +# Define task suite constants +class TaskSuite(str, Enum): + LIBERO_SPATIAL = "libero_spatial" + LIBERO_OBJECT = "libero_object" + LIBERO_GOAL = "libero_goal" + LIBERO_10 = "libero_10" + LIBERO_90 = "libero_90" + + +# Define max steps for each task suite +TASK_MAX_STEPS = { + TaskSuite.LIBERO_SPATIAL: 220, # longest training demo has 193 steps + TaskSuite.LIBERO_OBJECT: 280, # longest training demo has 254 steps + TaskSuite.LIBERO_GOAL: 300, # longest training demo has 270 steps + TaskSuite.LIBERO_10: 520, # longest training demo has 505 steps + TaskSuite.LIBERO_90: 400, # longest training demo has 373 steps +} + + +# Set up logging +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(message)s", + handlers=[logging.StreamHandler()], +) +logger = logging.getLogger(__name__) @dataclass @@ -60,72 +87,122 @@ class GenerateConfig: ################################################################################################################# model_family: str = "openvla" # Model family pretrained_checkpoint: Union[str, Path] = "" # Pretrained checkpoint path - load_in_8bit: bool = False # (For OpenVLA only) Load with 8-bit quantization - load_in_4bit: bool = False # (For OpenVLA only) Load with 4-bit quantization + + use_l1_regression: bool = True # If True, uses continuous action head with L1 regression objective + use_diffusion: bool = False # If True, uses continuous action head with diffusion modeling objective (DDIM) + num_diffusion_steps: int = 50 # (When `diffusion==True`) Number of diffusion steps for inference + use_film: bool = False # If True, uses FiLM to infuse language inputs into visual features + num_images_in_input: int = 2 # Number of images in the VLA input (default: 1) + use_proprio: bool = True # Whether to include proprio state in input center_crop: bool = True # Center crop? (if trained w/ random crop image aug) + num_open_loop_steps: int = 8 # Number of actions to execute open-loop before requerying policy + + unnorm_key: Union[str, Path] = "" # Action un-normalization key + + load_in_8bit: bool = False # (For OpenVLA only) Load with 8-bit quantization + load_in_4bit: bool = False # (For OpenVLA only) Load with 4-bit quantization ################################################################################################################# # LIBERO environment-specific parameters ################################################################################################################# - task_suite_name: str = "libero_spatial" # Task suite. Options: libero_spatial, libero_object, libero_goal, libero_10, libero_90 + task_suite_name: str = TaskSuite.LIBERO_SPATIAL # Task suite num_steps_wait: int = 10 # Number of steps to wait for objects to stabilize in sim num_trials_per_task: int = 50 # Number of rollouts per task + initial_states_path: str = "DEFAULT" # "DEFAULT", or path to initial states JSON file + env_img_res: int = 256 # Resolution for environment images (not policy input resolution) ################################################################################################################# # Utils ################################################################################################################# - run_id_note: Optional[str] = None # Extra note to add in run ID for logging + run_id_note: Optional[str] = None # Extra note to add to end of run ID for logging local_log_dir: str = "./experiments/logs" # Local directory for eval logs use_wandb: bool = False # Whether to also log results in Weights & Biases - wandb_project: str = "YOUR_WANDB_PROJECT" # Name of W&B project to log to (use default!) - wandb_entity: str = "YOUR_WANDB_ENTITY" # Name of entity to log under + wandb_entity: str = "your-wandb-entity" # Name of WandB entity + wandb_project: str = "your-wandb-project" # Name of WandB project seed: int = 7 # Random Seed (for reproducibility) # fmt: on -@draccus.wrap() -def eval_libero(cfg: GenerateConfig) -> None: - assert cfg.pretrained_checkpoint is not None, "cfg.pretrained_checkpoint must not be None!" - if "image_aug" in cfg.pretrained_checkpoint: +def validate_config(cfg: GenerateConfig) -> None: + """Validate configuration parameters.""" + assert cfg.pretrained_checkpoint is not None, "pretrained_checkpoint must not be None!" + + if "image_aug" in str(cfg.pretrained_checkpoint): assert cfg.center_crop, "Expecting `center_crop==True` because model was trained with image augmentations!" + assert not (cfg.load_in_8bit and cfg.load_in_4bit), "Cannot use both 8-bit and 4-bit quantization!" - # Set random seed - set_seed_everywhere(cfg.seed) + # Validate task suite + assert cfg.task_suite_name in [suite.value for suite in TaskSuite], f"Invalid task suite: {cfg.task_suite_name}" - # [OpenVLA] Set action un-normalization key - cfg.unnorm_key = cfg.task_suite_name +def initialize_model(cfg: GenerateConfig): + """Initialize model and associated components.""" # Load model model = get_model(cfg) - # [OpenVLA] Check that the model contains the action un-normalization key - if cfg.model_family == "openvla": - # In some cases, the key must be manually modified (e.g. after training on a modified version of the dataset - # with the suffix "_no_noops" in the dataset name) - if cfg.unnorm_key not in model.norm_stats and f"{cfg.unnorm_key}_no_noops" in model.norm_stats: - cfg.unnorm_key = f"{cfg.unnorm_key}_no_noops" - assert cfg.unnorm_key in model.norm_stats, f"Action un-norm key {cfg.unnorm_key} not found in VLA `norm_stats`!" + # Load proprio projector if needed + proprio_projector = None + if cfg.use_proprio: + proprio_projector = get_proprio_projector( + cfg, + model.llm_dim, + proprio_dim=8, # 8-dimensional proprio for LIBERO + ) + + # Load action head if needed + action_head = None + if cfg.use_l1_regression or cfg.use_diffusion: + action_head = get_action_head(cfg, model.llm_dim) - # [OpenVLA] Get Hugging Face processor + # Load noisy action projector if using diffusion + noisy_action_projector = None + if cfg.use_diffusion: + noisy_action_projector = get_noisy_action_projector(cfg, model.llm_dim) + + # Get OpenVLA processor if needed processor = None if cfg.model_family == "openvla": processor = get_processor(cfg) + check_unnorm_key(cfg, model) + + return model, action_head, proprio_projector, noisy_action_projector, processor + + +def check_unnorm_key(cfg: GenerateConfig, model) -> None: + """Check that the model contains the action un-normalization key.""" + # Initialize unnorm_key + unnorm_key = cfg.task_suite_name - # Initialize local logging + # In some cases, the key must be manually modified (e.g. after training on a modified version of the dataset + # with the suffix "_no_noops" in the dataset name) + if unnorm_key not in model.norm_stats and f"{unnorm_key}_no_noops" in model.norm_stats: + unnorm_key = f"{unnorm_key}_no_noops" + + assert unnorm_key in model.norm_stats, f"Action un-norm key {unnorm_key} not found in VLA `norm_stats`!" + + # Set the unnorm_key in cfg + cfg.unnorm_key = unnorm_key + + +def setup_logging(cfg: GenerateConfig): + """Set up logging to file and optionally to wandb.""" + # Create run ID run_id = f"EVAL-{cfg.task_suite_name}-{cfg.model_family}-{DATE_TIME}" if cfg.run_id_note is not None: run_id += f"--{cfg.run_id_note}" + + # Set up local logging os.makedirs(cfg.local_log_dir, exist_ok=True) local_log_filepath = os.path.join(cfg.local_log_dir, run_id + ".txt") log_file = open(local_log_filepath, "w") - print(f"Logging to local log file: {local_log_filepath}") + logger.info(f"Logging to local log file: {local_log_filepath}") - # Initialize Weights & Biases logging as well + # Initialize Weights & Biases logging if enabled if cfg.use_wandb: wandb.init( entity=cfg.wandb_entity, @@ -133,154 +210,319 @@ def eval_libero(cfg: GenerateConfig) -> None: name=run_id, ) + return log_file, local_log_filepath, run_id + + +def log_message(message: str, log_file=None): + """Log a message to console and optionally to a log file.""" + logger.info(message) + if log_file: + log_file.write(message + "\n") + log_file.flush() + + +def load_initial_states(cfg: GenerateConfig, task_suite, task_id: int, log_file=None): + """Load initial states for the given task.""" + # Get default initial states + initial_states = task_suite.get_task_init_states(task_id) + + # If using custom initial states, load them from file + if cfg.initial_states_path != "DEFAULT": + with open(cfg.initial_states_path, "r") as f: + all_initial_states = json.load(f) + log_message(f"Using initial states from {cfg.initial_states_path}", log_file) + return initial_states, all_initial_states + else: + log_message("Using default initial states", log_file) + return initial_states, None + + +def prepare_observation(obs, resize_size): + """Prepare observation for policy input.""" + # Get preprocessed images + img = get_libero_image(obs) + wrist_img = get_libero_wrist_image(obs) + + # Resize images to size expected by model + img_resized = resize_image_for_policy(img, resize_size) + wrist_img_resized = resize_image_for_policy(wrist_img, resize_size) + + # Prepare observations dict + observation = { + "full_image": img_resized, + "wrist_image": wrist_img_resized, + "state": np.concatenate( + (obs["robot0_eef_pos"], quat2axisangle(obs["robot0_eef_quat"]), obs["robot0_gripper_qpos"]) + ), + } + + return observation, img # Return both processed observation and original image for replay + + +def process_action(action, model_family): + """Process action before sending to environment.""" + # Normalize gripper action [0,1] -> [-1,+1] because the environment expects the latter + action = normalize_gripper_action(action, binarize=True) + + # [OpenVLA] The dataloader flips the sign of the gripper action to align with other datasets + # (0 = close, 1 = open), so flip it back (-1 = open, +1 = close) before executing the action + if model_family == "openvla": + action = invert_gripper_action(action) + + return action + + +def run_episode( + cfg: GenerateConfig, + env, + task_description: str, + model, + resize_size, + processor=None, + action_head=None, + proprio_projector=None, + noisy_action_projector=None, + initial_state=None, + log_file=None, +): + """Run a single episode in the environment.""" + # Reset environment + env.reset() + + # Set initial state if provided + if initial_state is not None: + obs = env.set_init_state(initial_state) + else: + obs = env.get_observation() + + # Initialize action queue + if cfg.num_open_loop_steps != NUM_ACTIONS_CHUNK: + print(f"WARNING: cfg.num_open_loop_steps ({cfg.num_open_loop_steps}) does not match the NUM_ACTIONS_CHUNK " + "{NUM_ACTIONS_CHUNK} constant defined in prismatic.vla.constants! For best performance (in terms of " + "both speed and success rate), we recommend executing the full action chunk.") + action_queue = deque(maxlen=cfg.num_open_loop_steps) + + # Setup + t = 0 + replay_images = [] + max_steps = TASK_MAX_STEPS[cfg.task_suite_name] + + # Run episode + success = False + try: + while t < max_steps + cfg.num_steps_wait: + # Do nothing for the first few timesteps to let objects stabilize + if t < cfg.num_steps_wait: + obs, reward, done, info = env.step(get_libero_dummy_action(cfg.model_family)) + t += 1 + continue + + # Prepare observation + observation, img = prepare_observation(obs, resize_size) + replay_images.append(img) + + # If action queue is empty, requery model + if len(action_queue) == 0: + # Query model to get action + actions = get_action( + cfg, + model, + observation, + task_description, + processor=processor, + action_head=action_head, + proprio_projector=proprio_projector, + noisy_action_projector=noisy_action_projector, + use_film=cfg.use_film, + ) + action_queue.extend(actions) + + # Get action from queue + action = action_queue.popleft() + + # Process action + action = process_action(action, cfg.model_family) + + # Execute action in environment + obs, reward, done, info = env.step(action.tolist()) + if done: + success = True + break + t += 1 + + except Exception as e: + log_message(f"Episode error: {e}", log_file) + + return success, replay_images + + +def run_task( + cfg: GenerateConfig, + task_suite, + task_id: int, + model, + resize_size, + processor=None, + action_head=None, + proprio_projector=None, + noisy_action_projector=None, + total_episodes=0, + total_successes=0, + log_file=None, +): + """Run evaluation for a single task.""" + # Get task + task = task_suite.get_task(task_id) + + # Get initial states + initial_states, all_initial_states = load_initial_states(cfg, task_suite, task_id, log_file) + + # Initialize environment and get task description + env, task_description = get_libero_env(task, cfg.model_family, resolution=cfg.env_img_res) + + # Start episodes + task_episodes, task_successes = 0, 0 + for episode_idx in tqdm.tqdm(range(cfg.num_trials_per_task)): + log_message(f"\nTask: {task_description}", log_file) + + # Handle initial state + if cfg.initial_states_path == "DEFAULT": + # Use default initial state + initial_state = initial_states[episode_idx] + else: + # Get keys for fetching initial episode state from JSON + initial_states_task_key = task_description.replace(" ", "_") + episode_key = f"demo_{episode_idx}" + + # Skip episode if expert demonstration failed to complete the task + if not all_initial_states[initial_states_task_key][episode_key]["success"]: + log_message(f"Skipping task {task_id} episode {episode_idx} due to failed expert demo!", log_file) + continue + + # Get initial state + initial_state = np.array(all_initial_states[initial_states_task_key][episode_key]["initial_state"]) + + log_message(f"Starting episode {task_episodes + 1}...", log_file) + + # Run episode + success, replay_images = run_episode( + cfg, + env, + task_description, + model, + resize_size, + processor, + action_head, + proprio_projector, + noisy_action_projector, + initial_state, + log_file, + ) + + # Update counters + task_episodes += 1 + total_episodes += 1 + if success: + task_successes += 1 + total_successes += 1 + + # Save replay video + save_rollout_video( + replay_images, total_episodes, success=success, task_description=task_description, log_file=log_file + ) + + # Log results + log_message(f"Success: {success}", log_file) + log_message(f"# episodes completed so far: {total_episodes}", log_file) + log_message(f"# successes: {total_successes} ({total_successes / total_episodes * 100:.1f}%)", log_file) + + # Log task results + task_success_rate = float(task_successes) / float(task_episodes) if task_episodes > 0 else 0 + total_success_rate = float(total_successes) / float(total_episodes) if total_episodes > 0 else 0 + + log_message(f"Current task success rate: {task_success_rate}", log_file) + log_message(f"Current total success rate: {total_success_rate}", log_file) + + # Log to wandb if enabled + if cfg.use_wandb: + wandb.log( + { + f"success_rate/{task_description}": task_success_rate, + f"num_episodes/{task_description}": task_episodes, + } + ) + + return total_episodes, total_successes + + +@draccus.wrap() +def eval_libero(cfg: GenerateConfig) -> float: + """Main function to evaluate a trained policy on LIBERO benchmark tasks.""" + # Validate configuration + validate_config(cfg) + + # Set random seed + set_seed_everywhere(cfg.seed) + + # Initialize model and components + model, action_head, proprio_projector, noisy_action_projector, processor = initialize_model(cfg) + + # Get expected image dimensions + resize_size = get_image_resize_size(cfg) + + # Setup logging + log_file, local_log_filepath, run_id = setup_logging(cfg) + # Initialize LIBERO task suite benchmark_dict = benchmark.get_benchmark_dict() task_suite = benchmark_dict[cfg.task_suite_name]() - num_tasks_in_suite = task_suite.n_tasks - print(f"Task suite: {cfg.task_suite_name}") - log_file.write(f"Task suite: {cfg.task_suite_name}\n") + num_tasks = task_suite.n_tasks - # Get expected image dimensions - resize_size = get_image_resize_size(cfg) + log_message(f"Task suite: {cfg.task_suite_name}", log_file) # Start evaluation total_episodes, total_successes = 0, 0 - for task_id in tqdm.tqdm(range(num_tasks_in_suite)): - # Get task - task = task_suite.get_task(task_id) - - # Get default LIBERO initial states - initial_states = task_suite.get_task_init_states(task_id) - - # Initialize LIBERO environment and task description - env, task_description = get_libero_env(task, cfg.model_family, resolution=256) - - # Start episodes - task_episodes, task_successes = 0, 0 - for episode_idx in tqdm.tqdm(range(cfg.num_trials_per_task)): - print(f"\nTask: {task_description}") - log_file.write(f"\nTask: {task_description}\n") - - # Reset environment - env.reset() - - # Set initial states - obs = env.set_init_state(initial_states[episode_idx]) - - # Setup - t = 0 - replay_images = [] - if cfg.task_suite_name == "libero_spatial": - max_steps = 220 # longest training demo has 193 steps - elif cfg.task_suite_name == "libero_object": - max_steps = 280 # longest training demo has 254 steps - elif cfg.task_suite_name == "libero_goal": - max_steps = 300 # longest training demo has 270 steps - elif cfg.task_suite_name == "libero_10": - max_steps = 520 # longest training demo has 505 steps - elif cfg.task_suite_name == "libero_90": - max_steps = 400 # longest training demo has 373 steps - - print(f"Starting episode {task_episodes+1}...") - log_file.write(f"Starting episode {task_episodes+1}...\n") - while t < max_steps + cfg.num_steps_wait: - try: - # IMPORTANT: Do nothing for the first few timesteps because the simulator drops objects - # and we need to wait for them to fall - if t < cfg.num_steps_wait: - obs, reward, done, info = env.step(get_libero_dummy_action(cfg.model_family)) - t += 1 - continue - - # Get preprocessed image - img = get_libero_image(obs, resize_size) - - # Save preprocessed image for replay video - replay_images.append(img) - - # Prepare observations dict - # Note: OpenVLA does not take proprio state as input - observation = { - "full_image": img, - "state": np.concatenate( - (obs["robot0_eef_pos"], quat2axisangle(obs["robot0_eef_quat"]), obs["robot0_gripper_qpos"]) - ), - } - - # Query model to get action - action = get_action( - cfg, - model, - observation, - task_description, - processor=processor, - ) - - # Normalize gripper action [0,1] -> [-1,+1] because the environment expects the latter - action = normalize_gripper_action(action, binarize=True) - - # [OpenVLA] The dataloader flips the sign of the gripper action to align with other datasets - # (0 = close, 1 = open), so flip it back (-1 = open, +1 = close) before executing the action - if cfg.model_family == "openvla": - action = invert_gripper_action(action) - - # Execute action in environment - obs, reward, done, info = env.step(action.tolist()) - if done: - task_successes += 1 - total_successes += 1 - break - t += 1 - - except Exception as e: - print(f"Caught exception: {e}") - log_file.write(f"Caught exception: {e}\n") - break - - task_episodes += 1 - total_episodes += 1 - - # Save a replay video of the episode - save_rollout_video( - replay_images, total_episodes, success=done, task_description=task_description, log_file=log_file - ) - - # Log current results - print(f"Success: {done}") - print(f"# episodes completed so far: {total_episodes}") - print(f"# successes: {total_successes} ({total_successes / total_episodes * 100:.1f}%)") - log_file.write(f"Success: {done}\n") - log_file.write(f"# episodes completed so far: {total_episodes}\n") - log_file.write(f"# successes: {total_successes} ({total_successes / total_episodes * 100:.1f}%)\n") - log_file.flush() - - # Log final results - print(f"Current task success rate: {float(task_successes) / float(task_episodes)}") - print(f"Current total success rate: {float(total_successes) / float(total_episodes)}") - log_file.write(f"Current task success rate: {float(task_successes) / float(task_episodes)}\n") - log_file.write(f"Current total success rate: {float(total_successes) / float(total_episodes)}\n") - log_file.flush() - if cfg.use_wandb: - wandb.log( - { - f"success_rate/{task_description}": float(task_successes) / float(task_episodes), - f"num_episodes/{task_description}": task_episodes, - } - ) - - # Save local log file - log_file.close() - - # Push total metrics and local log file to wandb + for task_id in tqdm.tqdm(range(num_tasks)): + total_episodes, total_successes = run_task( + cfg, + task_suite, + task_id, + model, + resize_size, + processor, + action_head, + proprio_projector, + noisy_action_projector, + total_episodes, + total_successes, + log_file, + ) + + # Calculate final success rate + final_success_rate = float(total_successes) / float(total_episodes) if total_episodes > 0 else 0 + + # Log final results + log_message("Final results:", log_file) + log_message(f"Total episodes: {total_episodes}", log_file) + log_message(f"Total successes: {total_successes}", log_file) + log_message(f"Overall success rate: {final_success_rate:.4f} ({final_success_rate * 100:.1f}%)", log_file) + + # Log to wandb if enabled if cfg.use_wandb: wandb.log( { - "success_rate/total": float(total_successes) / float(total_episodes), + "success_rate/total": final_success_rate, "num_episodes/total": total_episodes, } ) wandb.save(local_log_filepath) + # Close log file + if log_file: + log_file.close() + + return final_success_rate + if __name__ == "__main__": eval_libero() diff --git a/experiments/robot/libero/sample_libero_spatial_observation.pkl b/experiments/robot/libero/sample_libero_spatial_observation.pkl new file mode 100644 index 000000000..8863226be Binary files /dev/null and b/experiments/robot/libero/sample_libero_spatial_observation.pkl differ diff --git a/experiments/robot/openvla_utils.py b/experiments/robot/openvla_utils.py index e12e9a2f2..570d1cff5 100644 --- a/experiments/robot/openvla_utils.py +++ b/experiments/robot/openvla_utils.py @@ -1,48 +1,287 @@ -"""Utils for evaluating the OpenVLA policy.""" +"""Utils for evaluating OpenVLA or fine-tuned OpenVLA policies.""" +import filecmp import json import os +import shutil import time +from datetime import datetime +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union +import json_numpy import numpy as np +import requests import tensorflow as tf import torch +from huggingface_hub import HfApi, hf_hub_download from PIL import Image from transformers import AutoConfig, AutoImageProcessor, AutoModelForVision2Seq, AutoProcessor +# Apply JSON numpy patch for serialization +json_numpy.patch() + 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 +from prismatic.models.action_heads import DiffusionActionHead, L1RegressionActionHead +from prismatic.models.film_vit_wrapper import FiLMedPrismaticVisionBackbone +from prismatic.models.projectors import NoisyActionProjector, ProprioProjector +from prismatic.vla.constants import ( + ACTION_DIM, + ACTION_PROPRIO_NORMALIZATION_TYPE, +) +from prismatic.vla.datasets.rlds.utils.data_utils import NormalizationType -# Initialize important constants and pretty-printing mode in NumPy. -ACTION_DIM = 7 +# Initialize important constants DATE = time.strftime("%Y_%m_%d") DATE_TIME = time.strftime("%Y_%m_%d-%H_%M_%S") DEVICE = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") +OPENVLA_IMAGE_SIZE = 224 # Standard image size expected by OpenVLA + +# Configure NumPy print settings np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)}) -# Initialize system prompt for OpenVLA v0.1. -OPENVLA_V01_SYSTEM_PROMPT = ( - "A chat between a curious user and an artificial intelligence assistant. " - "The assistant gives helpful, detailed, and polite answers to the user's questions." -) +def model_is_on_hf_hub(model_path: str) -> bool: + """Checks whether a model path points to a model on Hugging Face Hub.""" + # If the API call below runs without error, the model is on the hub + try: + HfApi().model_info(model_path) + return True + except Exception: + return False + + +def update_auto_map(pretrained_checkpoint: str) -> None: + """ + Update the AutoMap configuration in the checkpoint config.json file. + + This loads the config.json file inside the checkpoint directory and overwrites + the AutoConfig and AutoModelForVision2Seq fields to use OpenVLA-specific classes. + + Args: + pretrained_checkpoint: Path to the checkpoint directory + """ + if not os.path.isdir(pretrained_checkpoint): + return + + config_path = os.path.join(pretrained_checkpoint, "config.json") + if not os.path.exists(config_path): + print(f"Warning: No config.json found at {config_path}") + return + + # Create timestamped backup + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + backup_path = os.path.join(pretrained_checkpoint, f"config.json.back.{timestamp}") + shutil.copy2(config_path, backup_path) + print(f"Created backup of original config at: {os.path.abspath(backup_path)}") + + # Read and update the config + with open(config_path, "r") as f: + config = json.load(f) + + config["auto_map"] = { + "AutoConfig": "configuration_prismatic.OpenVLAConfig", + "AutoModelForVision2Seq": "modeling_prismatic.OpenVLAForActionPrediction", + } + + # Write back the updated config + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + print(f"Updated config.json at: {os.path.abspath(config_path)}") + print("Changes made:") + print(' - Set AutoConfig to "configuration_prismatic.OpenVLAConfig"') + print(' - Set AutoModelForVision2Seq to "modeling_prismatic.OpenVLAForActionPrediction"') + + +def check_identical_files(path1: Union[str, Path], path2: Union[str, Path]) -> bool: + """ + Check if two files are identical in content. + + Args: + path1: Path to the first file + path2: Path to the second file + + Returns: + bool: True if files are identical, False otherwise + """ + path1, path2 = Path(path1), Path(path2) + + # First check if file sizes match + if path1.stat().st_size != path2.stat().st_size: + return False + + # Check if contents match + return filecmp.cmp(path1, path2, shallow=False) -def get_vla(cfg): - """Loads and returns a VLA model from checkpoint.""" - # Load VLA checkpoint. - print("[*] Instantiating Pretrained VLA model") - print("[*] Loading in BF16 with Flash-Attention Enabled") - # Register OpenVLA model to HF Auto Classes (not needed if the model is on HF Hub) - AutoConfig.register("openvla", OpenVLAConfig) - AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) - AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) - AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) +def _handle_file_sync(curr_filepath: str, checkpoint_filepath: str, file_type: str) -> None: + """ + Handle syncing of files between current directory and checkpoint. + + Creates backups if files exist but differ, and copies current versions to checkpoint. + + Args: + curr_filepath: Path to the current file version + checkpoint_filepath: Path where the file should be in the checkpoint + file_type: Description of the file type for logging + """ + if os.path.exists(checkpoint_filepath): + # Check if existing files are identical + match = check_identical_files(curr_filepath, checkpoint_filepath) + + if not match: + print( + "\n------------------------------------------------------------------------------------------------\n" + f"Found mismatch between:\n" + f"Current: {curr_filepath}\n" + f"Checkpoint: {checkpoint_filepath}\n" + ) + + # Create timestamped backup + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + backup_path = f"{checkpoint_filepath}.back.{timestamp}" + shutil.copy2(checkpoint_filepath, backup_path) + print(f"Created backup of original checkpoint file at: {os.path.abspath(backup_path)}") + + # Copy current version to checkpoint directory + shutil.copy2(curr_filepath, checkpoint_filepath) + print(f"Copied current version to checkpoint at: {os.path.abspath(checkpoint_filepath)}") + print( + f"Changes complete. The checkpoint will now use the current version of {file_type}" + "\n------------------------------------------------------------------------------------------------\n" + ) + else: + # If file doesn't exist in checkpoint directory, copy it + shutil.copy2(curr_filepath, checkpoint_filepath) + print( + "\n------------------------------------------------------------------------------------------------\n" + f"No {file_type} found in checkpoint directory.\n" + f"Copied current version from: {curr_filepath}\n" + f"To checkpoint location: {os.path.abspath(checkpoint_filepath)}" + "\n------------------------------------------------------------------------------------------------\n" + ) + + +def check_model_logic_mismatch(pretrained_checkpoint: str) -> None: + """ + Check and sync model logic files between current code and checkpoint. + + Handles the relationship between current and checkpoint versions of both + modeling_prismatic.py and configuration_prismatic.py: + - If checkpoint file exists and differs: creates backup and copies current version + - If checkpoint file doesn't exist: copies current version + + Args: + pretrained_checkpoint: Path to the checkpoint directory + """ + if not os.path.isdir(pretrained_checkpoint): + return + + # Find current files + curr_files = {"modeling_prismatic.py": None, "configuration_prismatic.py": None} + + for root, _, files in os.walk("./prismatic/"): + for filename in curr_files.keys(): + if filename in files and curr_files[filename] is None: + curr_files[filename] = os.path.join(root, filename) + + # Check and handle each file + for filename, curr_filepath in curr_files.items(): + if curr_filepath is None: + print(f"WARNING: `{filename}` is not found anywhere in the current directory.") + continue + + checkpoint_filepath = os.path.join(pretrained_checkpoint, filename) + _handle_file_sync(curr_filepath, checkpoint_filepath, filename) + + +def find_checkpoint_file(pretrained_checkpoint: str, file_pattern: str) -> str: + """ + Find a specific checkpoint file matching a pattern. + + Args: + pretrained_checkpoint: Path to the checkpoint directory + file_pattern: String pattern to match in filenames + + Returns: + str: Path to the matching checkpoint file + Raises: + AssertionError: If no files or multiple files match the pattern + """ + assert os.path.isdir(pretrained_checkpoint), f"Checkpoint path must be a directory: {pretrained_checkpoint}" + + checkpoint_files = [] + for filename in os.listdir(pretrained_checkpoint): + if file_pattern in filename and "checkpoint" in filename: + full_path = os.path.join(pretrained_checkpoint, filename) + checkpoint_files.append(full_path) + + assert len(checkpoint_files) == 1, ( + f"Expected exactly 1 {file_pattern} checkpoint but found {len(checkpoint_files)} in directory: {pretrained_checkpoint}" + ) + + return checkpoint_files[0] + + +def load_component_state_dict(checkpoint_path: str) -> Dict[str, torch.Tensor]: + """ + Load a component's state dict from checkpoint and handle DDP prefix if present. + + Args: + checkpoint_path: Path to the checkpoint file + + Returns: + Dict: The processed state dictionary for loading + """ + state_dict = torch.load(checkpoint_path, weights_only=True) + + # If the component was trained with DDP, elements in the state dict have prefix "module." which we must remove + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("module."): + new_state_dict[k[7:]] = v + else: + new_state_dict[k] = v + + return new_state_dict + + +def get_vla(cfg: Any) -> torch.nn.Module: + """ + Load and initialize the VLA model from checkpoint. + + Args: + cfg: Configuration object + + Returns: + torch.nn.Module: The initialized VLA model + """ + print("Instantiating pretrained VLA policy...") + + # If loading a locally stored pretrained checkpoint, check whether config or model files + # need to be synced so that any changes the user makes to the VLA modeling code will + # actually go into effect + # If loading a pretrained checkpoint from Hugging Face Hub, we just assume that the policy + # will be used as is, with its original modeling logic + if not model_is_on_hf_hub(cfg.pretrained_checkpoint): + # Register OpenVLA model to HF Auto Classes (not needed if the model is on HF Hub) + AutoConfig.register("openvla", OpenVLAConfig) + AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) + AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) + AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) + + # Update config.json and sync model files + update_auto_map(cfg.pretrained_checkpoint) + check_model_logic_mismatch(cfg.pretrained_checkpoint) + + # Load the model vla = AutoModelForVision2Seq.from_pretrained( cfg.pretrained_checkpoint, - attn_implementation="flash_attention_2", + # attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, load_in_8bit=cfg.load_in_8bit, load_in_4bit=cfg.load_in_4bit, @@ -50,14 +289,82 @@ def get_vla(cfg): trust_remote_code=True, ) - # Move model to device. - # Note: `.to()` is not supported for 8-bit or 4-bit bitsandbytes models, but the model will - # already be set to the right devices and casted to the correct dtype upon loading. + # If using FiLM, wrap the vision backbone to allow for infusion of language inputs + if cfg.use_film: + vla = _apply_film_to_vla(vla, cfg) + + # Set number of images in model input + vla.vision_backbone.set_num_images_in_input(cfg.num_images_in_input) + + vla.eval() + + # Move model to device if not using quantization if not cfg.load_in_8bit and not cfg.load_in_4bit: vla = vla.to(DEVICE) - # Load dataset stats used during finetuning (for action un-normalization). - dataset_statistics_path = os.path.join(cfg.pretrained_checkpoint, "dataset_statistics.json") + # Load dataset stats for action normalization + _load_dataset_stats(vla, cfg.pretrained_checkpoint) + + return vla + + +def _apply_film_to_vla(vla: torch.nn.Module, cfg: Any) -> torch.nn.Module: + """ + Apply FiLM (Feature-wise Linear Modulation) to the VLA vision backbone. + + Args: + vla: The VLA model + cfg: Configuration object with model parameters + + Returns: + torch.nn.Module: VLA model with FiLM applied + """ + from peft import LoraConfig, get_peft_model + + # Apply LoRA configuration + lora_config = LoraConfig( + r=32, + lora_alpha=16, + lora_dropout=0.0, + target_modules="all-linear", + init_lora_weights="gaussian", + ) + vla = get_peft_model(vla, lora_config) + + # Create and apply FiLMed vision backbone + new_vision_backbone = FiLMedPrismaticVisionBackbone( + vision_backbone=vla.vision_backbone, llm_dim=vla.llm_dim, + ) + vla.model.vision_backbone = new_vision_backbone + + # Load vision backbone checkpoint + checkpoint_path = find_checkpoint_file(cfg.pretrained_checkpoint, "vision_backbone") + state_dict = torch.load(checkpoint_path, weights_only=True) + vla.model.vision_backbone.load_state_dict(state_dict) + + # Use the model component instead of wrapper and convert to bfloat16 + vla = vla.model + vla.vision_backbone = vla.vision_backbone.to(torch.bfloat16) + + return vla + + +def _load_dataset_stats(vla: torch.nn.Module, checkpoint_path: str) -> None: + """ + Load dataset statistics used during training for action normalization. + + Args: + vla: The VLA model + checkpoint_path: Path to the checkpoint directory + """ + if model_is_on_hf_hub(checkpoint_path): + # Download dataset stats directly from HF Hub + dataset_statistics_path = hf_hub_download( + repo_id=checkpoint_path, + filename="dataset_statistics.json", + ) + else: + dataset_statistics_path = os.path.join(checkpoint_path, "dataset_statistics.json") if os.path.isfile(dataset_statistics_path): with open(dataset_statistics_path, "r") as f: norm_stats = json.load(f) @@ -69,39 +376,195 @@ def get_vla(cfg): "Otherwise, you may run into errors when trying to call `predict_action()` due to an absent `unnorm_key`." ) - return vla +def get_processor(cfg: Any) -> AutoProcessor: + """ + Get the VLA model's Hugging Face processor. -def get_processor(cfg): - """Get VLA model's Hugging Face processor.""" - processor = AutoProcessor.from_pretrained(cfg.pretrained_checkpoint, trust_remote_code=True) - return processor + Args: + cfg: Configuration object with model parameters + + Returns: + AutoProcessor: The model's processor + """ + return AutoProcessor.from_pretrained(cfg.pretrained_checkpoint, trust_remote_code=True) + + +def get_proprio_projector(cfg: Any, llm_dim: int, proprio_dim: int) -> ProprioProjector: + """ + Get proprioception projector for the VLA model. + + Args: + cfg: Configuration object with model parameters + llm_dim: Dimension of the language model + proprio_dim: Dimension of proprioception data + + Returns: + ProprioProjector: The initialized proprio projector + """ + # Initialize projector and move to device + proprio_projector = ProprioProjector( + llm_dim=llm_dim, + proprio_dim=proprio_dim, + ).to(DEVICE) + proprio_projector = proprio_projector.to(torch.bfloat16).to(DEVICE) + proprio_projector.eval() + + # Find and load checkpoint (may be on Hugging Face Hub or stored locally) + if model_is_on_hf_hub(cfg.pretrained_checkpoint): + model_path_to_proprio_projector_name = { + "moojink/openvla-7b-oft-finetuned-libero-spatial": "proprio_projector--150000_checkpoint.pt", + "moojink/openvla-7b-oft-finetuned-libero-object": "proprio_projector--150000_checkpoint.pt", + "moojink/openvla-7b-oft-finetuned-libero-goal": "proprio_projector--50000_checkpoint.pt", + "moojink/openvla-7b-oft-finetuned-libero-10": "proprio_projector--150000_checkpoint.pt", + } + if cfg.pretrained_checkpoint not in model_path_to_proprio_projector_name.keys(): + raise ValueError("Unsupported HF Hub pretrained checkpoint found!") + # Download proprio projector directly from HF Hub + proprio_projector_path = hf_hub_download( + repo_id=cfg.pretrained_checkpoint, filename=model_path_to_proprio_projector_name[cfg.pretrained_checkpoint] + ) + state_dict = load_component_state_dict(proprio_projector_path) + proprio_projector.load_state_dict(state_dict) + else: + checkpoint_path = find_checkpoint_file(cfg.pretrained_checkpoint, "proprio_projector") + state_dict = load_component_state_dict(checkpoint_path) + proprio_projector.load_state_dict(state_dict) + + return proprio_projector + + +def get_noisy_action_projector(cfg: Any, llm_dim: int) -> NoisyActionProjector: + """ + Get noisy action projector for diffusion-based action prediction. + + Args: + cfg: Configuration object with model parameters + llm_dim: Dimension of the language model + + Returns: + NoisyActionProjector: The initialized noisy action projector + """ + # Initialize projector and move to device + noisy_action_projector = NoisyActionProjector( + llm_dim=llm_dim, + ).to(DEVICE) + noisy_action_projector = noisy_action_projector.to(torch.bfloat16).to(DEVICE) + noisy_action_projector.eval() + + # Find and load checkpoint + checkpoint_path = find_checkpoint_file(cfg.pretrained_checkpoint, "noisy_action_projector") + state_dict = load_component_state_dict(checkpoint_path) + noisy_action_projector.load_state_dict(state_dict) + + return noisy_action_projector + + +def get_action_head(cfg: Any, llm_dim: int) -> Union[L1RegressionActionHead, DiffusionActionHead]: + """ + Get action head for continuous value prediction. + Args: + cfg: Configuration object with model parameters + llm_dim: Dimension of the language model + + Returns: + Union[L1RegressionActionHead, DiffusionActionHead]: The initialized action head -def crop_and_resize(image, crop_scale, batch_size): + Raises: + AssertionError: If both L1 regression and diffusion are specified """ - Center-crops an image to have area `crop_scale` * (original image area), and then resizes back - to original size. We use the same logic seen in the `dlimp` RLDS datasets wrapper to avoid - distribution shift at test time. + assert not (cfg.use_l1_regression and cfg.use_diffusion), "Cannot use both L1 regression and diffusion action head!" + + # Initialize appropriate action head based on configuration + if cfg.use_l1_regression: + action_head = L1RegressionActionHead(input_dim=llm_dim, hidden_dim=llm_dim, action_dim=ACTION_DIM) + elif cfg.use_diffusion: + action_head = DiffusionActionHead( + input_dim=llm_dim, hidden_dim=llm_dim, action_dim=ACTION_DIM, num_diffusion_steps=cfg.num_diffusion_steps + ) + else: + raise ValueError("Either use_l1_regression or use_diffusion must be True") + + action_head = action_head.to(torch.bfloat16).to(DEVICE) + action_head.eval() + + # Find and load checkpoint (may be on Hugging Face Hub or stored locally) + if model_is_on_hf_hub(cfg.pretrained_checkpoint): + model_path_to_action_head_name = { + "moojink/openvla-7b-oft-finetuned-libero-spatial": "action_head--150000_checkpoint.pt", + "moojink/openvla-7b-oft-finetuned-libero-object": "action_head--150000_checkpoint.pt", + "moojink/openvla-7b-oft-finetuned-libero-goal": "action_head--50000_checkpoint.pt", + "moojink/openvla-7b-oft-finetuned-libero-10": "action_head--150000_checkpoint.pt", + } + if cfg.pretrained_checkpoint not in model_path_to_action_head_name.keys(): + raise ValueError("Unsupported HF Hub pretrained checkpoint found!") + # Download proprio projector directly from HF Hub + action_head_path = hf_hub_download( + repo_id=cfg.pretrained_checkpoint, filename=model_path_to_action_head_name[cfg.pretrained_checkpoint] + ) + state_dict = load_component_state_dict(action_head_path) + action_head.load_state_dict(state_dict) + else: + checkpoint_path = find_checkpoint_file(cfg.pretrained_checkpoint, "action_head") + state_dict = load_component_state_dict(checkpoint_path) + action_head.load_state_dict(state_dict) + + return action_head + + +def resize_image_for_policy(img: np.ndarray, resize_size: Union[int, Tuple[int, int]]) -> np.ndarray: + """ + Resize an image to match the policy's expected input size. + + Uses the same resizing scheme as in the training data pipeline for distribution matching. Args: - image: TF Tensor of shape (batch_size, H, W, C) or (H, W, C) and datatype tf.float32 with - values between [0,1]. - crop_scale: The area of the center crop with respect to the original image. - batch_size: Batch size. + img: Numpy array containing the image + resize_size: Target size as int (square) or (height, width) tuple + + Returns: + np.ndarray: The resized image """ - # Convert from 3D Tensor (H, W, C) to 4D Tensor (batch_size, H, W, C) - assert image.shape.ndims == 3 or image.shape.ndims == 4 + assert isinstance(resize_size, int) or isinstance(resize_size, tuple) + if isinstance(resize_size, int): + resize_size = (resize_size, resize_size) + + # Resize using the same pipeline as in RLDS dataset builder + img = tf.image.encode_jpeg(img) # Encode as JPEG + img = tf.io.decode_image(img, expand_animations=False, dtype=tf.uint8) # Decode back + img = tf.image.resize(img, resize_size, method="lanczos3", antialias=True) + img = tf.cast(tf.clip_by_value(tf.round(img), 0, 255), tf.uint8) + + return img.numpy() + + +def crop_and_resize(image: tf.Tensor, crop_scale: float, batch_size: int) -> tf.Tensor: + """ + Center-crop an image and resize it back to original dimensions. + + Uses the same logic as in the training data pipeline for distribution matching. + + Args: + image: TF Tensor of shape (batch_size, H, W, C) or (H, W, C) with values in [0,1] + crop_scale: Area of center crop relative to original image + batch_size: Batch size + + Returns: + tf.Tensor: The cropped and resized image + """ + # Handle 3D inputs by adding batch dimension if needed + assert image.shape.ndims in (3, 4), "Image must be 3D or 4D tensor" expanded_dims = False if image.shape.ndims == 3: image = tf.expand_dims(image, axis=0) expanded_dims = True - # Get height and width of crop + # Calculate crop dimensions (note: we use sqrt(crop_scale) for h/w) new_heights = tf.reshape(tf.clip_by_value(tf.sqrt(crop_scale), 0, 1), shape=(batch_size,)) new_widths = tf.reshape(tf.clip_by_value(tf.sqrt(crop_scale), 0, 1), shape=(batch_size,)) - # Get bounding box representing crop + # Create bounding box for the crop height_offsets = (1 - new_heights) / 2 width_offsets = (1 - new_widths) / 2 bounding_boxes = tf.stack( @@ -114,57 +577,238 @@ def crop_and_resize(image, crop_scale, batch_size): axis=1, ) - # Crop and then resize back up - image = tf.image.crop_and_resize(image, bounding_boxes, tf.range(batch_size), (224, 224)) + # Apply crop and resize + image = tf.image.crop_and_resize( + image, bounding_boxes, tf.range(batch_size), (OPENVLA_IMAGE_SIZE, OPENVLA_IMAGE_SIZE) + ) - # Convert back to 3D Tensor (H, W, C) + # Remove batch dimension if it was added if expanded_dims: image = image[0] return image -def get_vla_action(vla, processor, base_vla_name, obs, task_label, unnorm_key, center_crop=False): - """Generates an action with the VLA policy.""" - image = Image.fromarray(obs["full_image"]) - image = image.convert("RGB") +def center_crop_image(image: Union[np.ndarray, Image.Image]) -> Image.Image: + """ + Center crop an image to match training data distribution. + + Args: + image: Input image (PIL or numpy array) - # (If trained with image augmentations) Center crop image and then resize back up to original size. - # IMPORTANT: Let's say crop scale == 0.9. To get the new height and width (post-crop), multiply - # the original height and width by sqrt(0.9) -- not 0.9! - if center_crop: - batch_size = 1 - crop_scale = 0.9 + Returns: + Image.Image: Cropped PIL Image + """ + batch_size = 1 + crop_scale = 0.9 - # Convert to TF Tensor and record original data type (should be tf.uint8) + # Convert to TF Tensor if needed + if not isinstance(image, tf.Tensor): image = tf.convert_to_tensor(np.array(image)) - orig_dtype = image.dtype - # Convert to data type tf.float32 and values between [0,1] - image = tf.image.convert_image_dtype(image, tf.float32) + orig_dtype = image.dtype - # Crop and then resize back to original size - image = crop_and_resize(image, crop_scale, batch_size) + # Convert to float32 in range [0,1] + image = tf.image.convert_image_dtype(image, tf.float32) - # Convert back to original data type - image = tf.clip_by_value(image, 0, 1) - image = tf.image.convert_image_dtype(image, orig_dtype, saturate=True) + # Apply center crop and resize + image = crop_and_resize(image, crop_scale, batch_size) - # Convert back to PIL Image - image = Image.fromarray(image.numpy()) - image = image.convert("RGB") + # Convert back to original data type + image = tf.clip_by_value(image, 0, 1) + image = tf.image.convert_image_dtype(image, orig_dtype, saturate=True) - # Build VLA prompt - if "openvla-v01" in base_vla_name: # OpenVLA v0.1 - prompt = ( - f"{OPENVLA_V01_SYSTEM_PROMPT} USER: What action should the robot take to {task_label.lower()}? ASSISTANT:" - ) - else: # OpenVLA + # Convert to PIL Image + return Image.fromarray(image.numpy()).convert("RGB") + + +def check_image_format(image: Any) -> None: + """ + Validate input image format. + + Args: + image: Image to check + + Raises: + AssertionError: If image format is invalid + """ + is_numpy_array = isinstance(image, np.ndarray) + has_correct_shape = len(image.shape) == 3 and image.shape[-1] == 3 + has_correct_dtype = image.dtype == np.uint8 + + assert is_numpy_array and has_correct_shape and has_correct_dtype, ( + "Incorrect image format detected! Make sure that the input image is a " + "numpy array with shape (H, W, 3) and dtype np.uint8!" + ) + + +def normalize_proprio(proprio: np.ndarray, norm_stats: Dict[str, Any]) -> np.ndarray: + """ + Normalize proprioception data to match training distribution. + + Args: + proprio: Raw proprioception data + norm_stats: Normalization statistics + + Returns: + np.ndarray: Normalized proprioception data + """ + if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS: + mask = norm_stats.get("mask", np.ones_like(norm_stats["min"], dtype=bool)) + proprio_high, proprio_low = np.array(norm_stats["max"]), np.array(norm_stats["min"]) + elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99: + mask = norm_stats.get("mask", np.ones_like(norm_stats["q01"], dtype=bool)) + proprio_high, proprio_low = np.array(norm_stats["q99"]), np.array(norm_stats["q01"]) + else: + raise ValueError("Unsupported action/proprio normalization type detected!") + + normalized_proprio = np.clip( + np.where( + mask, + 2 * (proprio - proprio_low) / (proprio_high - proprio_low + 1e-8) - 1, + proprio, + ), + a_min=-1.0, + a_max=1.0, + ) + + return normalized_proprio + + +def prepare_images_for_vla(images: List[np.ndarray], cfg: Any) -> List[Image.Image]: + """ + Prepare images for VLA input by resizing and cropping as needed. + + Args: + images: List of input images as numpy arrays + cfg: Configuration object with parameters + + Returns: + List[Image.Image]: Processed images ready for the model + """ + processed_images = [] + + for image in images: + # Validate format + check_image_format(image) + + # Resize if needed + if image.shape != (OPENVLA_IMAGE_SIZE, OPENVLA_IMAGE_SIZE, 3): + image = resize_image_for_policy(image, OPENVLA_IMAGE_SIZE) + + # Convert to PIL image + pil_image = Image.fromarray(image).convert("RGB") + + # Apply center crop if configured + if cfg.center_crop: + pil_image = center_crop_image(pil_image) + + processed_images.append(pil_image) + + return processed_images + + +def get_vla_action( + cfg: Any, + vla: torch.nn.Module, + processor: Any, + obs: Dict[str, Any], + task_label: str, + action_head: Optional[torch.nn.Module] = None, + proprio_projector: Optional[torch.nn.Module] = None, + noisy_action_projector: Optional[torch.nn.Module] = None, + use_film: bool = False, +) -> List[np.ndarray]: + """ + Generate action predictions with the VLA policy. + + Args: + cfg: Configuration object with parameters + vla: The VLA model + processor: Model processor for inputs + obs: Observation dictionary + task_label: Text description of the task + action_head: Optional action head for continuous actions + proprio_projector: Optional proprioception projector + noisy_action_projector: Optional noisy action projector for diffusion + use_film: Whether to use FiLM + + Returns: + List[np.ndarray]: Predicted actions + """ + with torch.inference_mode(): + + # Collect all input images + all_images = [obs["full_image"]] + if cfg.num_images_in_input > 1: + all_images.extend([obs[k] for k in obs.keys() if "wrist" in k]) + + # Process images + all_images = prepare_images_for_vla(all_images, cfg) + + # Extract primary image and additional images + primary_image = all_images.pop(0) + + # Build VLA prompt prompt = f"In: What action should the robot take to {task_label.lower()}?\nOut:" - # Process inputs. - inputs = processor(prompt, image).to(DEVICE, dtype=torch.bfloat16) + # Process primary image + inputs = processor(prompt, primary_image).to(DEVICE, dtype=torch.bfloat16) + + # Process additional wrist images if any + if all_images: + all_wrist_inputs = [ + processor(prompt, image_wrist).to(DEVICE, dtype=torch.bfloat16) for image_wrist in all_images + ] + # Concatenate all images + primary_pixel_values = inputs["pixel_values"] + all_wrist_pixel_values = [wrist_inputs["pixel_values"] for wrist_inputs in all_wrist_inputs] + inputs["pixel_values"] = torch.cat([primary_pixel_values] + all_wrist_pixel_values, dim=1) + + # Process proprioception data if used + proprio = None + if cfg.use_proprio: + proprio = obs["state"] + proprio_norm_stats = vla.norm_stats[cfg.unnorm_key]["proprio"] + obs["state"] = normalize_proprio(proprio, proprio_norm_stats) + proprio = obs["state"] + + # Generate action + if action_head is None: + # Standard VLA output (single-image inputs, discrete actions) + action, _ = vla.predict_action(**inputs, unnorm_key=cfg.unnorm_key, do_sample=False) + else: + # Custom action head for continuous actions + action, _ = vla.predict_action( + **inputs, + unnorm_key=cfg.unnorm_key, + do_sample=False, + proprio=proprio, + proprio_projector=proprio_projector, + noisy_action_projector=noisy_action_projector, + action_head=action_head, + use_film=use_film, + ) + + # Extract subset of actions for open loop steps + return [action[i] for i in range(min(len(action), cfg.num_open_loop_steps))] + + +def get_action_from_server( + observation: Dict[str, Any], server_endpoint: str = "http://0.0.0.0:8777/act" +) -> Dict[str, Any]: + """ + Get VLA action from remote inference server. + + Args: + observation: Observation data to send to server + server_endpoint: URL of the inference server - # Get action. - action = vla.predict_action(**inputs, unnorm_key=unnorm_key, do_sample=False) - return action + Returns: + Dict[str, Any]: Action response from server + """ + response = requests.post( + server_endpoint, + json=observation, + ) + return response.json() diff --git a/experiments/robot/robot_utils.py b/experiments/robot/robot_utils.py index 10e5289d8..64559e990 100644 --- a/experiments/robot/robot_utils.py +++ b/experiments/robot/robot_utils.py @@ -3,6 +3,7 @@ import os import random import time +from typing import Any, Dict, List, Optional, Union import numpy as np import torch @@ -12,22 +13,35 @@ get_vla_action, ) -# Initialize important constants and pretty-printing mode in NumPy. +# Initialize important constants ACTION_DIM = 7 DATE = time.strftime("%Y_%m_%d") DATE_TIME = time.strftime("%Y_%m_%d-%H_%M_%S") DEVICE = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") + +# Configure NumPy print settings np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)}) -# Initialize system prompt for OpenVLA v0.1. +# Initialize system prompt for OpenVLA v0.1 OPENVLA_V01_SYSTEM_PROMPT = ( "A chat between a curious user and an artificial intelligence assistant. " "The assistant gives helpful, detailed, and polite answers to the user's questions." ) +# Model image size configuration +MODEL_IMAGE_SIZES = { + "openvla": 224, + # Add other models as needed +} + -def set_seed_everywhere(seed: int): - """Sets the random seed for Python, NumPy, and PyTorch functions.""" +def set_seed_everywhere(seed: int) -> None: + """ + Set random seed for all random number generators for reproducibility. + + Args: + seed: The random seed to use + """ torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) @@ -37,66 +51,149 @@ def set_seed_everywhere(seed: int): os.environ["PYTHONHASHSEED"] = str(seed) -def get_model(cfg, wrap_diffusion_policy_for_droid=False): - """Load model for evaluation.""" +def get_model(cfg: Any, wrap_diffusion_policy_for_droid: bool = False) -> torch.nn.Module: + """ + Load and initialize model for evaluation based on configuration. + + Args: + cfg: Configuration object with model parameters + wrap_diffusion_policy_for_droid: Whether to wrap diffusion policy for DROID + + Returns: + torch.nn.Module: The loaded model + + Raises: + ValueError: If model family is not supported + """ if cfg.model_family == "openvla": model = get_vla(cfg) else: - raise ValueError("Unexpected `model_family` found in config.") + raise ValueError(f"Unsupported model family: {cfg.model_family}") + print(f"Loaded model: {type(model)}") return model -def get_image_resize_size(cfg): +def get_image_resize_size(cfg: Any) -> Union[int, tuple]: """ - Gets image resize size for a model class. - If `resize_size` is an int, then the resized image will be a square. - Else, the image will be a rectangle. - """ - if cfg.model_family == "openvla": - resize_size = 224 - else: - raise ValueError("Unexpected `model_family` found in config.") - return resize_size + Get image resize dimensions for a specific model. + If returned value is an int, the resized image will be a square. + If returned value is a tuple, the resized image will be a rectangle. + + Args: + cfg: Configuration object with model parameters + + Returns: + Union[int, tuple]: Image resize dimensions + + Raises: + ValueError: If model family is not supported + """ + if cfg.model_family not in MODEL_IMAGE_SIZES: + raise ValueError(f"Unsupported model family: {cfg.model_family}") + + return MODEL_IMAGE_SIZES[cfg.model_family] + + +def get_action( + cfg: Any, + model: torch.nn.Module, + obs: Dict[str, Any], + task_label: str, + processor: Optional[Any] = None, + action_head: Optional[torch.nn.Module] = None, + proprio_projector: Optional[torch.nn.Module] = None, + noisy_action_projector: Optional[torch.nn.Module] = None, + use_film: bool = False, +) -> Union[List[np.ndarray], np.ndarray]: + """ + Query the model to get action predictions. + + Args: + cfg: Configuration object with model parameters + model: The loaded model + obs: Observation dictionary + task_label: Text description of the task + processor: Model processor for inputs + action_head: Optional action head for continuous actions + proprio_projector: Optional proprioception projector + noisy_action_projector: Optional noisy action projector for diffusion + use_film: Whether to use FiLM + + Returns: + Union[List[np.ndarray], np.ndarray]: Predicted actions + + Raises: + ValueError: If model family is not supported + """ + with torch.no_grad(): + if cfg.model_family == "openvla": + action = get_vla_action( + cfg=cfg, + vla=model, + processor=processor, + obs=obs, + task_label=task_label, + action_head=action_head, + proprio_projector=proprio_projector, + noisy_action_projector=noisy_action_projector, + use_film=use_film, + ) + else: + raise ValueError(f"Unsupported model family: {cfg.model_family}") -def get_action(cfg, model, obs, task_label, processor=None): - """Queries the model to get an action.""" - if cfg.model_family == "openvla": - action = get_vla_action( - model, processor, cfg.pretrained_checkpoint, obs, task_label, cfg.unnorm_key, center_crop=cfg.center_crop - ) - assert action.shape == (ACTION_DIM,) - else: - raise ValueError("Unexpected `model_family` found in config.") return action -def normalize_gripper_action(action, binarize=True): +def normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray: """ - Changes gripper action (last dimension of action vector) from [0,1] to [-1,+1]. - Necessary for some environments (not Bridge) because the dataset wrapper standardizes gripper actions to [0,1]. - Note that unlike the other action dimensions, the gripper action is not normalized to [-1,+1] by default by - the dataset wrapper. + Normalize gripper action from [0,1] to [-1,+1] range. + + This is necessary for some environments because the dataset wrapper + standardizes gripper actions to [0,1]. Note that unlike the other action + dimensions, the gripper action is not normalized to [-1,+1] by default. Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1 + + Args: + action: Action array with gripper action in the last dimension + binarize: Whether to binarize gripper action to -1 or +1 + + Returns: + np.ndarray: Action array with normalized gripper action """ - # Just normalize the last action to [-1,+1]. + # Create a copy to avoid modifying the original + normalized_action = action.copy() + + # Normalize the last action dimension to [-1,+1] orig_low, orig_high = 0.0, 1.0 - action[..., -1] = 2 * (action[..., -1] - orig_low) / (orig_high - orig_low) - 1 + normalized_action[..., -1] = 2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1 if binarize: - # Binarize to -1 or +1. - action[..., -1] = np.sign(action[..., -1]) + # Binarize to -1 or +1 + normalized_action[..., -1] = np.sign(normalized_action[..., -1]) - return action + return normalized_action -def invert_gripper_action(action): +def invert_gripper_action(action: np.ndarray) -> np.ndarray: """ - Flips the sign of the gripper action (last dimension of action vector). - This is necessary for some environments where -1 = open, +1 = close, since + Flip the sign of the gripper action (last dimension of action vector). + + This is necessary for environments where -1 = open, +1 = close, since the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open. + + Args: + action: Action array with gripper action in the last dimension + + Returns: + np.ndarray: Action array with inverted gripper action """ - action[..., -1] = action[..., -1] * -1.0 - return action + # Create a copy to avoid modifying the original + inverted_action = action.copy() + + # Invert the gripper action + inverted_action[..., -1] *= -1.0 + + return inverted_action diff --git a/finetune_lerobot.sh b/finetune_lerobot.sh new file mode 100755 index 000000000..9ded1dc8c --- /dev/null +++ b/finetune_lerobot.sh @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +PYTHONPATH="$PYTHONPATH:$(pwd)/lerobot" && \ +torchrun --standalone --nnodes 1 --nproc-per-node 1 vla-scripts/finetune.py \ + --vla_path "openvla/openvla-7b" \ + --data_root_dir "data/robotgeneralist" \ + --dataset_name "nomagic-simple-box" \ + --run_root_dir ".runs/" \ + --shuffle_buffer_size 100000 \ + --use_l1_regression True \ + --use_diffusion False \ + --num_diffusion_steps 50 \ + --use_film True \ + --num_images_in_input 1 \ + --use_proprio False \ + --batch_size 1 \ + --learning_rate 5e-4 \ + --lr_warmup_steps 0 \ + --num_steps_before_decay 100000 \ + --grad_accumulation_steps 8 \ + --max_steps 200000 \ + --use_val_set True \ + --val_freq 10000 \ + --val_time_limit 180 \ + --save_freq 1000 \ + --save_latest_checkpoint_only False \ + --resume False \ + --resume_step None \ + --image_aug True \ + --diffusion_sample_freq 50 \ + --use_lora True \ + --lora_rank 32 \ + --lora_dropout 0.1 \ + --merge_lora_during_training False \ + --wandb_entity robotgeneralist \ + --wandb_project ur5e \ + --wandb_log_freq 10 \ + --use_lerobot_dataset True \ + --lerobot_dataset_root_dir "data" \ + --lerobot_dataset_name "robotgeneralist/nomagic-simple-box" \ + --lerobot_tolerance_s 0.01 + diff --git a/manual_create_env.sh b/manual_create_env.sh new file mode 100755 index 000000000..333f06396 --- /dev/null +++ b/manual_create_env.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# +# Creates a Python 3.10 virtual environment, installs dependencies, +# and downloads the OpenVLA model. Intended for environments +# without Docker support where GPU is Slurm-managed (e.g. entropy). +# Designed to be used in conjunction with vla-scripts/finetune.sub. +# +# Usage: ./manual_create_env.sh +# +# Assumptions: +# - This script is run from the root directory of the openvla_finetuner project. +# - The user has a working Python 3.10 and `virtualenv` installation. +# - The user has internet access to download dependencies and the model. +# - The user has Slurm installed and configured. + +set -e +set -o pipefail + +function check_virtualenv { + if ! command -v virtualenv &> /dev/null; then + echo "Error: virtualenv is not installed" >&2 + exit 1 + fi +} + +function verify_directory { + if [[ ! -f "$(basename "$0")" ]]; then + echo "Error: This script must be run from the openvla/ directory" >&2 + exit 1 + fi +} + +function setup_virtualenv { + if [[ ! -d ".venv" ]]; then + virtualenv -p 3.10 .venv || exit 1 + fi + source .venv/bin/activate || exit 1 + pip install --upgrade pip || exit 1 + pip install "setuptools<60" || exit 1 # Or else build dlimp_openvla will fail. +} + +function install_dependencies { + # Note: We intentionally don't exit on pip check error below due to known + # dependency conflicts between OpenVLA (which needs torch==2.2.0) and LeRobot + # (which needs torch>=2.2.1). These conflicts are expected and the + # installation will still work for our purposes. + + # Install OpenVLA. + pip install -e . || exit 1 + + # Install LeRobot. + pushd third_party/lerobot || exit 1 + pip install -e . || exit 1 + popd || exit 1 + + # Reinstall OpenVLA dependencies that LeRobot may have overwritten. + set +e # Temporarily disable exit on error. + pip check | awk '$1 ~ /openvla/ {gsub(/,/,"",$5); print $5}' | \ + xargs pip install + set -e # Re-enable exit on error. + + # Install Flash Attention. + pip install packaging ninja || exit 1 + # We will install flash-attn inside the Slurm job, after loading CUDA. + # This avoids errors when CUDA_HOME is not set during environment creation. + # pip install \ + # "flash-attn==2.5.5" \ + # --no-build-isolation || exit 1 +} + +function download_model { + pip install huggingface-hub || exit 1 + huggingface-cli download openvla/openvla-7b || exit 1 +} + +function create_save_dirs { + verify_directory + mkdir -p data + mkdir -p .runs + mkdir -p .slurmlog +} + +function main { + check_virtualenv + verify_directory + setup_virtualenv + install_dependencies + download_model + create_save_dirs +} + +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + main "$@" +fi \ No newline at end of file diff --git a/prismatic/extern/hf/modeling_prismatic.py b/prismatic/extern/hf/modeling_prismatic.py index 4a26c4871..9013a4f57 100644 --- a/prismatic/extern/hf/modeling_prismatic.py +++ b/prismatic/extern/hf/modeling_prismatic.py @@ -1,15 +1,9 @@ """ modeling_prismatic.py -Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions, inheriting -from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained, but exactly replicate the -logic in `prismatic.models.vlms.prismatic.py`. - -Note =>> for the time being, not adding the custom HF "docstring" formatting. - -References [LLaVa, IDEFICS-2]: - => https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/modeling_llava.py - => https://github.com/huggingface/transformers/blob/main/src/transformers/models/idefics2/modeling_idefics2.py +Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions. +Inherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained, +but exactly replicate the logic in `prismatic.models.vlms.prismatic.py`. """ import logging @@ -27,16 +21,26 @@ from transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import ModelOutput +from prismatic.training.train_utils import ( + get_current_action_mask, + get_next_actions_mask, +) +from prismatic.vla.constants import ( + ACTION_DIM, + ACTION_PROPRIO_NORMALIZATION_TYPE, + ACTION_TOKEN_BEGIN_IDX, + IGNORE_INDEX, + NUM_ACTIONS_CHUNK, + STOP_INDEX, + NormalizationType, +) + from .configuration_prismatic import OpenVLAConfig, PrismaticConfig -# Get Logger +# Set up logger logger = logging.getLogger(__name__) -# === PyTorch/HuggingFace Default IGNORE_INDEX (for CrossEntropyLoss labels) -IGNORE_INDEX = -100 - - # === Utility Functions for Monkey-Patching === def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]: def wrapper(*args: Any, **kwargs: Any) -> Any: @@ -61,6 +65,13 @@ def ls_apply_patch(ls_module: LayerScale): # === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) === class PrismaticVisionBackbone(nn.Module): + """ + Vision backbone for Prismatic models that handles image feature extraction. + + Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations. + For fused backbones, features from both models are concatenated along the feature dimension. + """ + def __init__( self, use_fused_vision_backbone: bool, @@ -68,59 +79,152 @@ def __init__( timm_model_ids: List[str], timm_override_act_layers: List[Optional[str]], ) -> None: + """ + Initialize the vision backbone. + + Args: + use_fused_vision_backbone: Whether to use two backbones and fuse their features + image_sizes: List of image sizes for each backbone + timm_model_ids: List of TIMM model IDs to use for each backbone + timm_override_act_layers: List of activation layer overrides for each backbone + """ super().__init__() self.use_fused_vision_backbone = use_fused_vision_backbone + self.num_images_in_input = 1 # Default value, can be overridden later - # [Contract] Validate number of (fused) vision backbones, create "alpha" featurizer and Instantiate - # =>> Note :: Monkey-Patch the `forward()` function of the backbone to ensure FSDP-compatibility - # Hardcodes `get_intermediate_layers` to return the **SECOND-TO-LAST** layer patches! - assert len(timm_model_ids) <= 2, "Prismatic models only support up to 2 (fused) vision backbones!" - self.featurizer = timm.create_model( - timm_model_ids[0], - pretrained=False, - num_classes=0, - img_size=image_sizes[0], - act_layer=timm_override_act_layers[0], - ) - self.featurizer.forward = unpack_tuple( - partial(self.featurizer.get_intermediate_layers, n={len(self.featurizer.blocks) - 2}) + # Validate number of (fused) vision backbones + if len(timm_model_ids) > 2: + raise ValueError("Prismatic models only support up to 2 (fused) vision backbones!") + + # Create primary featurizer + self.featurizer = self._create_featurizer( + model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0] ) self.embed_dim = self.featurizer.embed_dim - # If `use_fused_vision_backbone` =>> create "beta" featurizer + # Create secondary featurizer if using fused backbone if self.use_fused_vision_backbone: - self.fused_featurizer = timm.create_model( - timm_model_ids[1], - pretrained=False, - num_classes=0, - img_size=image_sizes[1], - act_layer=timm_override_act_layers[1], - ) - self.fused_featurizer.forward = unpack_tuple( - partial(self.fused_featurizer.get_intermediate_layers, n={len(self.fused_featurizer.blocks) - 2}) + self.fused_featurizer = self._create_featurizer( + model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1] ) self.embed_dim += self.fused_featurizer.embed_dim - # Patch `vision_backbone.featurizer` and `vision_backbone.fused_featurizer` with HF-Compatible LayerScale + # Patch LayerScale modules for HF compatibility + self._patch_layer_scales() + + def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module: + """ + Create a TIMM-based featurizer model with appropriate configurations. + + Args: + model_id: The TIMM model ID to load + img_size: Input image size for the model + act_layer: Override for the activation layer type + + Returns: + A configured featurizer model + """ + featurizer = timm.create_model( + model_id, + pretrained=False, + num_classes=0, + img_size=img_size, + act_layer=act_layer, + ) + + # Monkey-patch the forward function to extract the second-to-last layer features + num_blocks = len(featurizer.blocks) + featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2})) + + return featurizer + + def _patch_layer_scales(self) -> None: + """ + Patch all LayerScale modules to be compatible with HF's parameter naming. + + HF Transformers overwrites parameters with names containing 'gamma', + so we need to rename and modify the forward method. + """ + # Patch primary featurizer for module in self.featurizer.modules(): if isinstance(module, LayerScale): ls_apply_patch(module) + # Patch secondary featurizer if it exists if self.use_fused_vision_backbone: for module in self.fused_featurizer.modules(): if isinstance(module, LayerScale): ls_apply_patch(module) + def get_num_patches(self) -> int: + """ + Returns the number of vision patches output by the vision backbone. + + Returns: + Number of patches per image + """ + return self.featurizer.patch_embed.num_patches + + def get_num_images_in_input(self) -> int: + """ + Returns the number of input images for the vision backbone. + + Returns: + Number of images expected in the input + """ + return self.num_images_in_input + + def set_num_images_in_input(self, num_images_in_input: int) -> None: + """ + Sets the number of input images for the vision backbone. + + Args: + num_images_in_input: Number of images to expect in the input + """ + self.num_images_in_input = num_images_in_input + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - """Run image (`pixel_values`) through featurizer; if channel-stacked, then dispatch and sequence stack.""" - if not self.use_fused_vision_backbone: - return self.featurizer(pixel_values) + """ + Implements the forward pass for the vision backbone. + + If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features + (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone). + + Args: + pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W). + """ + if self.num_images_in_input == 1: + if not self.use_fused_vision_backbone: + return self.featurizer(pixel_values) + + # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack + img, img_fused = torch.split(pixel_values, [3, 3], dim=1) + patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused) + + return torch.cat([patches, patches_fused], dim=2) + + else: + assert self.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!" + + # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2) + images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1) - # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack - img, img_fused = torch.split(pixel_values, [3, 3], dim=1) - patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused) + # Process each image and collect patches + all_patches = [] + for img in images: + # Split each image further into two stacks of channels (each with 3 channels) + img_regular, img_fused = torch.split(img, [3, 3], dim=1) - return torch.cat([patches, patches_fused], dim=2) + # Get patches from both SigLIP and DINOv2 vision transformers + patches = self.featurizer(img_regular) + patches_fused = self.fused_featurizer(img_fused) + + # Concatenate SigLIP and DINOv2 patches along the hidden dimension + combined_patches = torch.cat([patches, patches_fused], dim=2) + all_patches.append(combined_patches) + + # Concatenate all patches along the patch dimension + return torch.cat(all_patches, dim=1) # === Prismatic Projector (nn.Module) Definitions === @@ -250,6 +354,7 @@ def __init__(self, config: PrismaticConfig) -> None: ) self.vocab_size = config.text_config.vocab_size self.pad_token_id = config.pad_token_id + self.llm_dim = config.text_config.hidden_size # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing self.post_init() @@ -287,6 +392,109 @@ def resize_token_embeddings( return updated_embeddings + def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features): + """ + Replace embeddings in input_embeddings at positions where all_actions_mask is True + with embeddings from noisy_action_features, using vectorized operations. + + Args: + input_embeddings: Tensor of shape (B, S, D) + all_actions_mask: Boolean tensor of shape (B, S) + noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample + + Returns: + Modified input_embeddings tensor + """ + # Clone input to avoid modifying the original tensor + new_input_embeddings = input_embeddings.clone() + + # Create a tensor with the same shape of input_embeddings to hold the noisy action features + repositioned_noisy_action_features = torch.zeros_like(input_embeddings) + + # Create batch indices for splicing + batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device) + batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1]) + + # Get indices where mask is True for each sample + masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask]) + + # Move the noisy action features into their correct positions + repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features + + # Combine original input embeddings and noisy action embeddings using the mask + new_input_embeddings = torch.where( + all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings + ) + + return new_input_embeddings + + def _process_action_masks(self, labels): + """Helper to get action masks from labels""" + current_action_mask = get_current_action_mask(labels) + next_actions_mask = get_next_actions_mask(labels) + all_actions_mask = current_action_mask | next_actions_mask # (B, seq_len) + return all_actions_mask + + def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False): + """Process vision features with optional FiLM conditioning""" + if use_film: + # FiLM: Infuse language inputs into visual features + patch_features = self.vision_backbone(pixel_values, language_embeddings) # (bsz, 256 * num_images, D) + else: + patch_features = self.vision_backbone(pixel_values) # (bsz, 256 * num_images, D) + + # Project patch embeddings into language embedding space + return self.projector(patch_features) + + def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector): + """Process proprioceptive features and append to vision features""" + if proprio_projector is not None and proprio is not None: + # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim) + # proprio: (bsz, proprio_dim) or (propro_dim,) + proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1) # (bsz, proprio_dim) + proprio_features = proprio_projector(proprio) # (bsz, llm_dim) + proprio_features = proprio_features.unsqueeze(dim=1) # (bsz, 1, llm_dim) + # For simplicity, just append proprio token to the end of projected vision patch tokens + return torch.cat((projected_patch_embeddings, proprio_features), dim=1) + return projected_patch_embeddings + + def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask): + """Build multimodal embeddings and attention mask""" + # Update attention mask + projected_patch_attention_mask = None + if attention_mask is not None: + projected_patch_attention_mask = torch.full( + (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]), + fill_value=True, + dtype=attention_mask.dtype, + device=attention_mask.device, + ) + + # Build multimodal embeddings & attention mask; insert embeddings after token (1:) + multimodal_embeddings = torch.cat( + [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1 + ) + + multimodal_attention_mask = None + if attention_mask is not None: + multimodal_attention_mask = torch.cat( + [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1 + ) + + return multimodal_embeddings, multimodal_attention_mask + + def _build_multimodal_labels(self, labels, projected_patch_embeddings): + """Build multimodal labels with IGNORE_INDEX for patch embeddings""" + if labels is not None: + projected_patch_labels = torch.full( + (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]), + fill_value=IGNORE_INDEX, + dtype=labels.dtype, + device=labels.device, + ) + return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) + return None + # === Core Prismatic VLM `forward()` Logic === def forward( self, @@ -301,6 +509,12 @@ def forward( output_hidden_states: Optional[bool] = None, output_projector_features: Optional[bool] = None, return_dict: Optional[bool] = None, + proprio=None, + proprio_projector=None, + noisy_actions=None, + noisy_action_projector=None, + diffusion_timestep_embeddings=None, + use_film: bool = False, ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]: """Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance.""" output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions @@ -316,11 +530,6 @@ def forward( # Instantiate Placeholder for Projector Features projected_patch_embeddings = None - # Note :: We only support forward passes with the following cases: - # => Cached Generation :: (input_ids.shape[1] == 1) and (past_key_values is not None) - # => Unimodal Forward :: (pixel_values is None) - # => Multimodal Forward :: (pixel_values is not None) and (input_ids/embeds.shape[0] == pixel_values.shape[0]) - # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` === if input_ids.shape[1] == 1: assert input_ids.shape[0] == 1, "Generation is only currently supported for batch size of 1!" @@ -360,47 +569,66 @@ def forward( # === Handle Multimodal Forward === elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]): - assert past_key_values is None, "Unexpected key `past_key_values` provided during language-only forward!" + assert past_key_values is None, "Unexpected key `past_key_values` provided during multimodal forward!" - # Visual Feature Extraction - patch_features = self.vision_backbone(pixel_values) - - # Projection Logic =>> Update Attention Mask - projected_patch_embeddings = self.projector(patch_features) - projected_patch_attention_mask = None - if attention_mask is not None: - projected_patch_attention_mask = torch.full( - (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]), - fill_value=True, - dtype=attention_mask.dtype, - device=attention_mask.device, - ) + # Get input embeddings (from language model embeddings) + input_embeddings = self.get_input_embeddings()(input_ids) # (B, seq_len, D) + + # Extract action masks + all_actions_mask = self._process_action_masks(labels) - # Get Input Embeddings (from Language Model Embeddings) - input_embeddings = self.get_input_embeddings()(input_ids) + # Extract the language portion of the input embeddings (i.e. remove the action tokens portion) + language_embeddings = input_embeddings[~all_actions_mask].reshape( + input_embeddings.shape[0], -1, input_embeddings.shape[2] + ) # (B, lang_seq_len, llm_dim) - # Build Multimodal Embeddings & Attention Mask =>> Prismatic defaults to inserting after token (1:) - multimodal_embeddings = torch.cat( - [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1 + # Get visual features + projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film) + + # Add proprioceptive state if provided + projected_patch_embeddings = self._process_proprio_features( + projected_patch_embeddings, proprio, proprio_projector ) - multimodal_attention_mask = None - if attention_mask is not None: - multimodal_attention_mask = torch.cat( - [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1 + + # [Diffusion] Add diffusion timestep embedding if provided + if diffusion_timestep_embeddings is not None: + # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens + projected_patch_embeddings = torch.cat( + (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1 ) - # Build Labels (if specified) =>> Ignore Labels for Patch Embeddings - multimodal_labels = None - if labels is not None: - projected_patch_labels = torch.full( - (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]), - fill_value=IGNORE_INDEX, - dtype=labels.dtype, - device=labels.device, + # Process action embeddings + if noisy_actions is not None: + # Get mask corresponding to all action tokens + all_actions_mask = self._process_action_masks(labels) + + # Reshape noisy actions into individual action tokens + # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1) + B = noisy_actions.shape[0] + noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1) + + # Project noisy action tokens into language model embedding space + noisy_action_features = noisy_action_projector(noisy_actions) # (B, chunk_len * action_dim, llm_dim) + + # Replace embeddings of the action tokens with noisy action embeddings + input_embeddings = self._replace_input_embeddings( + input_embeddings, all_actions_mask, noisy_action_features ) - multimodal_labels = torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1) + else: + # Replace the embeddings of the action tokens with zeros + # (Later on, the positional embeddings will be added to them) + all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1) + input_embeddings = input_embeddings * ~all_actions_mask + + # Build multimodal embeddings & attention mask + multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention( + input_embeddings, projected_patch_embeddings, attention_mask + ) + + # Build labels for multimodal sequence if needed + multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings) - # Dispatch to Language Model + # Dispatch to language model language_model_output = self.language_model( input_ids=None, attention_mask=multimodal_attention_mask, @@ -503,10 +731,244 @@ def __init__(self, config: OpenVLAConfig) -> None: # Compute vocab size for de-tokenization -- revert added "multiple of" self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of + def _prepare_input_for_action_prediction(self, input_ids, attention_mask): + """Prepares input for action prediction by adding necessary tokens""" + # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens + placeholder_action_token_ids = ( + torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype) + ) + input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1) + + # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time) + stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX + input_ids = torch.cat([input_ids, stop_token_id], dim=-1) + + # Extend the attention mask to fit the new shape of input + # Note: Only batch size == 1 supported right now + mask_extension = ( + torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1])) + .to(attention_mask.device) + .to(attention_mask.dtype) + ) + attention_mask = torch.cat([attention_mask, mask_extension], dim=-1) + + return input_ids, attention_mask + + def _prepare_labels_for_action_prediction(self, labels, input_ids): + """Creates labels tensor for action prediction if not provided""" + # Extend labels tensor with fake action labels + ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1 + labels_extension = ( + torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype) + * ARBITRARY_ACTION_TOKEN_IDX + ) + labels = torch.cat([labels, labels_extension], dim=-1) + + # Replace last label token with stop token + labels[:, -1] = STOP_INDEX + + return labels + + def _unnormalize_actions(self, normalized_actions, unnorm_key=None): + """Unnormalize actions using dataset statistics""" + action_norm_stats = self.get_action_stats(unnorm_key) + + if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS: + mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["min"], dtype=bool)) + action_high, action_low = np.array(action_norm_stats["max"]), np.array(action_norm_stats["min"]) + elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99: + mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool)) + action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"]) + else: + raise ValueError("Unsupported action/proprio normalization type detected!") + + actions = np.where( + mask, + 0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low, + normalized_actions, + ) + + return actions + + def _run_diffusion_prediction( + self, + input_embeddings, + all_actions_mask, + noise, + action_head, + projected_patch_embeddings, + labels, + attention_mask, + NUM_PATCHES, + NUM_PROMPT_TOKENS, + noisy_action_projector, + ): + """Run diffusion-based action prediction""" + # Set diffusion timestep values + action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps) + # Clone embedding for reuse in each timestep + orig_projected_patch_embeddings = projected_patch_embeddings.clone() + curr_noisy_actions = noise + + # Reverse diffusion: Iteratively denoise to generate action prediction + for t in action_head.noise_scheduler.timesteps: + # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action + # embedding, and diffusion timestep embedding) + timesteps = torch.Tensor([t]).to(labels.device) + diffusion_timestep_embeddings = ( + action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device) + ) # (B, llm_dim) + diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim) + + # [Diffusion] Replace the embeddings of the action tokens with noisy actions + # (Later on, the positional embeddings will be added to them) + + # For simplicity, append diffusion timestep embedding to the end of projected vision tokens + projected_patch_embeddings = torch.cat( + (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1 + ) + + # Reshape and project noisy actions into language embedding space + B = curr_noisy_actions.shape[0] + orig_curr_noisy_actions_shape = curr_noisy_actions.shape + curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1) + noisy_action_features = noisy_action_projector(curr_noisy_actions) + curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape) + + # Replace action token embeddings with noisy action embeddings + input_embeddings = self._replace_input_embeddings( + input_embeddings.clone(), all_actions_mask, noisy_action_features + ) + + # Build multimodal embeddings and attention mask + multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention( + input_embeddings, projected_patch_embeddings, attention_mask + ) + + # Forward pass through language model + language_model_output = self.language_model( + input_ids=None, + attention_mask=multimodal_attention_mask, + position_ids=None, + past_key_values=None, + inputs_embeds=multimodal_embeddings, + labels=None, + use_cache=None, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + + # Extract hidden states for action portion of response + last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D) + actions_hidden_states = last_hidden_states[ + :, + NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK, + :, + ] # (B, act_chunk_len, D) + + # Predict noise and update noisy actions: x_t -> x_{t-1} + noise_pred = action_head.predict_noise(actions_hidden_states) + curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample + + curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM) + + # Return final actions + return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states + + def _regression_or_discrete_prediction( + self, + input_embeddings, + all_actions_mask, + projected_patch_embeddings, + attention_mask, + labels, + NUM_PATCHES, + NUM_PROMPT_TOKENS, + action_head=None, + ): + """Run L1 regression-based continuous action prediction or discrete action tokens prediction.""" + # Zero out action token embeddings + all_actions_mask = all_actions_mask.unsqueeze(-1) # (B, seq_len, 1) + input_embeddings = input_embeddings * ~all_actions_mask + + # Build multimodal embeddings and attention mask + multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention( + input_embeddings, projected_patch_embeddings, attention_mask + ) + + # Forward pass through language model + language_model_output = self.language_model( + input_ids=None, + attention_mask=multimodal_attention_mask, + position_ids=None, + past_key_values=None, + inputs_embeds=multimodal_embeddings, + labels=None, + use_cache=None, + output_attentions=False, + output_hidden_states=True, + return_dict=True, + ) + + # Extract hidden states for action tokens + last_hidden_states = language_model_output.hidden_states[-1] # (B, seq_len, D) + actions_hidden_states = last_hidden_states[ + :, + NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK, + :, + ] # (B, act_chunk_len, D) + + # Handle different prediction methods + if action_head is not None: + # L1 regression prediction + normalized_actions = action_head.predict_action(actions_hidden_states) + normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM) + normalized_actions = normalized_actions.float().cpu().detach().numpy() + else: + # Discrete token-based prediction + predicted_action_token_ids = ( + language_model_output.logits[ + :, + NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK, + ] + .argmax(dim=2) + .cpu() + .numpy() + ) + discretized_actions = self.vocab_size - predicted_action_token_ids + discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1) + normalized_actions = self.bin_centers[discretized_actions] + normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM) + + return normalized_actions, actions_hidden_states + def predict_action( - self, input_ids: Optional[torch.LongTensor] = None, unnorm_key: Optional[str] = None, **kwargs: str + self, + input_ids: Optional[torch.LongTensor] = None, + unnorm_key: Optional[str] = None, + proprio=None, + proprio_projector=None, + action_head=None, + noisy_action_projector=None, + use_film: bool = False, + **kwargs: str, ) -> np.ndarray: - """Thin wrapper around .generate() that decodes predicted actions and unnormalizes them.""" + """Predict actions from input sequence, with options for different prediction methods. + + Args: + input_ids: Input token ids + unnorm_key: Key for unnormalization statistics + proprio: Proprioceptive features + proprio_projector: Projector for proprioceptive features + action_head: Optional head for L1 regression or diffusion-based prediction + noisy_action_projector: Projector for noisy actions in diffusion-based prediction + use_film: Whether to use FiLM conditioning + **kwargs: Additional arguments including pixel_values and attention_mask + + Returns: + Tuple of (unnormalized_actions, action_hidden_states) + """ # If the special empty token ('') does not already appear after the colon (':') token in the prompt # (after "OUT:" or "ASSISTANT:"), insert it to match the inputs seen at training time if not torch.all(input_ids[:, -1] == 29871): @@ -514,29 +976,92 @@ def predict_action( (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1 ) - # Run VLA inference - generated_ids = self.generate(input_ids, max_new_tokens=self.get_action_dim(unnorm_key), **kwargs) + pixel_values = kwargs["pixel_values"] + attention_mask = kwargs["attention_mask"] - # Extract predicted action tokens and translate into (normalized) continuous actions - predicted_action_token_ids = generated_ids[0, -self.get_action_dim(unnorm_key) :].cpu().numpy() - discretized_actions = self.vocab_size - predicted_action_token_ids - discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1) - normalized_actions = self.bin_centers[discretized_actions] + # Create fake labels tensor (needed for action mask) + labels = input_ids.clone() + labels[:] = IGNORE_INDEX - # Unnormalize actions - action_norm_stats = self.get_action_stats(unnorm_key) - mask = action_norm_stats.get("mask", np.ones_like(action_norm_stats["q01"], dtype=bool)) - action_high, action_low = np.array(action_norm_stats["q99"]), np.array(action_norm_stats["q01"]) - actions = np.where( - mask, - 0.5 * (normalized_actions + 1) * (action_high - action_low) + action_low, - normalized_actions, + # Get number of tokens in prompt (excluding the start token) + NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1 # Subtract action tokens and stop token + + # Prepare inputs by adding necessary tokens + input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask) + + # Update labels tensor for action mask computation later + labels = self._prepare_labels_for_action_prediction(labels, input_ids) + + # Get input embeddings and action masks + input_embeddings = self.get_input_embeddings()(input_ids) + all_actions_mask = self._process_action_masks(labels) + + # Extract language embeddings + language_embeddings = input_embeddings[~all_actions_mask].reshape( + input_embeddings.shape[0], -1, input_embeddings.shape[2] ) - return actions + # Process vision features + projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film) + + # Add proprioceptive features if provided + use_proprio = proprio_projector is not None and proprio is not None + if use_proprio: + proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype) + projected_patch_embeddings = self._process_proprio_features( + projected_patch_embeddings, proprio, proprio_projector + ) + + # Use diffusion if provided, otherwise use regression or discrete prediction + use_diffusion = noisy_action_projector is not None and hasattr(action_head, "noise_scheduler") + + # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present) + NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input() + if use_proprio: + NUM_PATCHES += 1 + if use_diffusion: + NUM_PATCHES += 1 + + if use_diffusion: + # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion + noise = torch.randn( + size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype + ) + + # Run diffusion-based prediction + normalized_actions, actions_hidden_states = self._run_diffusion_prediction( + input_embeddings, + all_actions_mask, + noise, + action_head, + projected_patch_embeddings, + labels, + attention_mask, + NUM_PATCHES, + NUM_PROMPT_TOKENS, + noisy_action_projector, + ) + else: + # Run regression or discrete token-based prediction + normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction( + input_embeddings, + all_actions_mask, + projected_patch_embeddings, + attention_mask, + labels, + NUM_PATCHES, + NUM_PROMPT_TOKENS, + action_head, + ) + + # Unnormalize predicted actions + actions = self._unnormalize_actions(normalized_actions, unnorm_key) + + return actions, actions_hidden_states @staticmethod def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str: + """Validate and resolve the unnormalization key for action statistics""" if unnorm_key is None: assert len(norm_stats) == 1, ( f"Your model was trained on more than one dataset, " @@ -554,7 +1079,7 @@ def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optiona def get_action_dim(self, unnorm_key: Optional[str] = None) -> int: """Get the dimensionality of the policy's action space.""" unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key) - return len(self.norm_stats[unnorm_key]["action"]["q01"]) + return len(self.norm_stats[unnorm_key]["action"]["min"]) def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]: """Get all the logged statistics for the given dataset.""" diff --git a/prismatic/models/action_heads.py b/prismatic/models/action_heads.py new file mode 100644 index 000000000..b3043c078 --- /dev/null +++ b/prismatic/models/action_heads.py @@ -0,0 +1,211 @@ +"""Implementations of various action heads, which serve as alternatives to VLM sequential token prediction.""" + +import math + +import numpy as np +import torch +import torch.nn as nn +from diffusers.schedulers.scheduling_ddim import DDIMScheduler +from prismatic.vla.constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX, NUM_ACTIONS_CHUNK, PROPRIO_DIM, STOP_INDEX + + +class SinusoidalPositionalEncoding(nn.Module): + """ + Sine- and cosine-based positional encoding that produces embeddings of a batch of timesteps. + + For example, at train time, the input might be a batch of 32 randomly sampled diffusion timesteps -> shape (32,) + Then the output would be a batch of 32 timestep embeddings -> shape (32, D) + + Adapted from: https://github.com/real-stanford/diffusion_policy/blob/main/diffusion_policy/model/diffusion/positional_embedding.py + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim # dimensionality of the positional encoding + + def forward(self, x): + # x: (batch_size,) + device = x.device + assert self.dim % 2 == 0, f"# dimensions must be even but got {self.dim}" + half_dim = self.dim // 2 + exponent = torch.arange(half_dim, device=device) * -math.log(10000) / (half_dim - 1) # shape: (D/2,) + emb = torch.exp(exponent) # shape: (D/2,) + emb = x[:, None] * emb[None, :] # shape: (batch_size, 1) * (1, D/2) -> (batch_size, D/2) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) # shape: (batch_size, D) + return emb + + +class MLPResNetBlock(nn.Module): + """One MLP ResNet block with a residual connection.""" + def __init__(self, dim): + super().__init__() + self.dim = dim + self.ffn = nn.Sequential( # feedforward network, similar to the ones in Transformers + nn.LayerNorm(dim), + nn.Linear(dim, dim), + nn.ReLU(), + ) + + def forward(self, x): + # x: (batch_size, hidden_dim) + # We follow the module ordering of "Pre-Layer Normalization" feedforward networks in Transformers as + # described here: https://arxiv.org/pdf/2002.04745.pdf + identity = x + x = self.ffn(x) + x = x + identity + return x + + +class MLPResNet(nn.Module): + """MLP with residual connection blocks.""" + def __init__(self, num_blocks, input_dim, hidden_dim, output_dim): + super().__init__() + self.layer_norm1 = nn.LayerNorm(input_dim) + self.fc1 = nn.Linear(input_dim, hidden_dim) + self.relu = nn.ReLU() + self.mlp_resnet_blocks = nn.ModuleList() + for _ in range(num_blocks): + self.mlp_resnet_blocks.append(MLPResNetBlock(dim=hidden_dim)) + self.layer_norm2 = nn.LayerNorm(hidden_dim) + self.fc2 = nn.Linear(hidden_dim, output_dim) + + def forward(self, x): + # x: (batch_size, input_dim) + x = self.layer_norm1(x) # shape: (batch_size, input_dim) + x = self.fc1(x) # shape: (batch_size, hidden_dim) + x = self.relu(x) # shape: (batch_size, hidden_dim) + for block in self.mlp_resnet_blocks: + x = block(x) # shape: (batch_size, hidden_dim) + x = self.layer_norm2(x) # shape: (batch_size, hidden_dim) + x = self.fc2(x) # shape: (batch_size, output_dim) + return x + + +class L1RegressionActionHead(nn.Module): + """Simple MLP-based action head that generates continuous actions via L1 regression.""" + def __init__( + self, + input_dim=4096, + hidden_dim=4096, + action_dim=7, + ): + super().__init__() + self.action_dim = action_dim + self.model = MLPResNet( + num_blocks=2, input_dim=input_dim*ACTION_DIM, hidden_dim=hidden_dim, output_dim=action_dim + ) + + def predict_action(self, actions_hidden_states): + # actions_hidden_states: last hidden states of Transformer corresponding to action tokens in sequence + # - shape: (batch_size, chunk_len * action_dim, hidden_dim) + # ground_truth_actions: ground-truth actions + # - shape: (batch_size, chunk_len, action_dim) + batch_size = actions_hidden_states.shape[0] + device = actions_hidden_states.device + rearranged_actions_hidden_states = actions_hidden_states.reshape(batch_size, NUM_ACTIONS_CHUNK, -1) + action = self.model(rearranged_actions_hidden_states) + return action + + +class NoisePredictionModel(nn.Module): + """ + Diffusion noise prediction model that takes an observation embedding (which fuses the + noisy action, diffusion timestep, and image-language observation embeddings) and + outputs a noise prediction. + """ + + def __init__( + self, + transformer_hidden_dim, # Transformer hidden embedding size + hidden_dim, # MLP hidden size + action_dim=7, # action dimensionality + ): + super().__init__() + self.mlp_resnet = MLPResNet( + num_blocks=2, + input_dim=transformer_hidden_dim, + hidden_dim=hidden_dim, + output_dim=action_dim, + ) + + def forward( + self, + obs, + ): + # obs: observation embeddings to condition the generation on + # - shape: (batch_size, chunk_len, rearranged_hidden_dim=action_dim*hidden_dim) + # + # output: predicted noise + # - shape: (batch_size, action_dim) + output = self.mlp_resnet(obs) + return output + + +class DiffusionActionHead(nn.Module): + """ + Simple MLP-based action head that generates continuous actions via conditional denoising diffusion process. + + Loosely inspired by: https://github.com/real-stanford/diffusion_policy/blob/main/diffusion_policy/model/diffusion/transformer_for_diffusion.py + """ + + def __init__( + self, + input_dim=4096, + hidden_dim=4096, + action_dim=7, + num_diffusion_steps=100, + ): + super().__init__() + self.action_dim = action_dim + self.noise_predictor = NoisePredictionModel( + transformer_hidden_dim=hidden_dim*ACTION_DIM, hidden_dim=hidden_dim, action_dim=action_dim + ) + self.noise_scheduler = DDIMScheduler(num_train_timesteps=num_diffusion_steps, beta_schedule="squaredcos_cap_v2") + self.num_diffusion_steps = num_diffusion_steps + self.time_encoder = SinusoidalPositionalEncoding(dim=hidden_dim) + + def sample_noisy_actions(self, ground_truth_actions): + """ + Samples noise and applies noise to ground-truth actions to produce noisy actions, which are + used as input in the noise prediction network. Returns noise, noisy actions, and the + corresponding diffusion timestep embeddings. + """ + # ground_truth_actions: ground-truth actions + # - shape: (batch_size, chunk_len, action_dim) + batch_size = ground_truth_actions.shape[0] + device = ground_truth_actions.device + # Sample random noise with shape equal to actions, used for closed-form forward diffusion. + noise = torch.randn(size=(batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM), device=device, dtype=ground_truth_actions.dtype) # (B, chunk_len, action_dim) + # Sample random diffusion timesteps (one for each action in batch). + timesteps = torch.randint( + low=0, high=self.noise_scheduler.config.num_train_timesteps, size=(batch_size,), device=device + ) + # Add noise to clean actions according to the magnitude at each diffusion timestep via + # closed-form forward diffusion. + noisy_actions = self.noise_scheduler.add_noise(ground_truth_actions, noise, timesteps) # (B, chunk_len, action_dim) + + # Get diffusion timestep embeddings as well + diffusion_timestep_embeddings = self.time_encoder(timesteps).to(noisy_actions.dtype).to(noisy_actions.device) # (B, llm_dim) + diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim) + + return_dict = dict( + noise=noise, + noisy_actions=noisy_actions, + diffusion_timestep_embeddings=diffusion_timestep_embeddings, + ) + + return return_dict + + def predict_noise(self, actions_hidden_states): + """ + Given a batch of last hidden Transformer layer embeddings (which fuse the vision-language observation embeddings, + noisy action embeddings, and diffusion timestep embedding), predicts the noise applied to the actions. + """ + # actions_hidden_states: last hidden states of Transformer corresponding to action tokens in sequence + # - shape: (batch_size, chunk_len * action_dim, hidden_dim) + batch_size = actions_hidden_states.shape[0] + device = actions_hidden_states.device + rearranged_actions_hidden_states = actions_hidden_states.reshape(batch_size, NUM_ACTIONS_CHUNK, -1) # (batch_size, chunk_len, action_dim * hidden_dim) + # Get diffusion model's noise prediction. + noise_pred = self.noise_predictor(rearranged_actions_hidden_states) + return noise_pred diff --git a/prismatic/models/film_vit_wrapper.py b/prismatic/models/film_vit_wrapper.py new file mode 100644 index 000000000..94618ca96 --- /dev/null +++ b/prismatic/models/film_vit_wrapper.py @@ -0,0 +1,276 @@ +"""Implementation of additional modules for the VLA's vision transformer.""" + +from functools import partial +from typing import Any, Callable, Sequence, Tuple, Union + +import torch +import torch.nn as nn +from timm.models.vision_transformer import VisionTransformer + + +class FiLMedVisionTransformerBlock(nn.Module): + """ + Wrapper for ViT blocks that adds components to implement FiLM language conditioning. + + Modulates visual feature embeddings via + x = (1 + gamma) * x + beta, + where x is visual feature and gamma and beta are learned projections of the average language embedding. + gamma and beta have D dimensions each, where D is the number of hidden dimensions in the ViT's features. + + NOTE #1 (Moo Jin): + In convolutional neural architectures, the "feature" in FiLM is an entire feature map, i.e., each channel in a + convolutional layer (so gamma and beta have C dimensions, where C is the number of channels). Therefore, FiLM's + scaling and shifting is applied across all spatial locations for conv nets -- i.e., it is spatially agnostic. + + For vision transformer architectures, you may consider individual patch embeddings as individual "features" at first + instinct, but this would make FiLM scaling and shifting spatially local. In order to make the modulation spatially + global like in convolutional architectures, we should apply the scaling and shifting to each dimension of each patch + embedding. I.e., gamma and beta should have D dimensions, where D is the number of dimensions in a visual embedding. + + NOTE #2 (Moo Jin): + x = (1 + gamma) * x + beta is used in the original FiLM paper as opposed to x = gamma * x + beta (see section 7.2 in + https://arxiv.org/pdf/1709.07871.pdf). Since gamma and beta are close to zero upon initialization, this leads to an + identity transformation at the beginning of training, which minimizes perturbation to the pretrained representation. + """ + + def __init__( + self, + block, + vision_dim: int, + llm_dim: int, + ): + """ + Initializes FiLM ViT block wrapper. + + Args: + block (timm.models.vision_transformer.Block): Vision transformer block. + vision_dim (int): Number of hidden dimensions in visual embeddings. + llm_dim (int): Number of hidden dimensions in language embeddings. + """ + super().__init__() + self.block = block + # Initialize gamma and beta projectors + self.scale = nn.Linear(llm_dim, vision_dim) + self.shift = nn.Linear(llm_dim, vision_dim) + + def forward(self, x, average_language_embedding): + """ + Overrides the vision transformer block forward pass to use FiLM. + + Args: + x (torch.Tensor): Visual input embeddings, (batch_size, vision_seq_len, vision_dim). + average_language_embedding (torch.Tensor): Average language embedding for task, (batch_size, llm_dim). + """ + # Project average language embedding to visual embedding space to get gamma and beta + gamma = self.scale(average_language_embedding) # (batch_size, vision_dim) + beta = self.shift(average_language_embedding) # (batch_size, vision_dim) + + # Pass visual inputs through attention portion of original block + x = x + self.block.drop_path1(self.block.ls1(self.block.attn(self.block.norm1(x)))) + + # Modulate intermediate visual representations via FiLM + x = x * (1 + gamma.view(gamma.shape[0], 1, gamma.shape[1])) + beta.view(beta.shape[0], 1, beta.shape[1]) + + # Pass visual inputs through feedforward portion of original block + x = x + self.block.drop_path2(self.block.ls2(self.block.mlp(self.block.norm2(x)))) + + return x + + +class NullVisionTransformerBlockWrapper(nn.Module): + """ + Null wrapper for ViT blocks that doesn't do anything; just calls the original block's forward function. + Useful if you want to use a block wrapper every X blocks instead of every block (e.g., to reduce the number of new + parameters introduced by a new wrapper). + """ + + def __init__( + self, + block, + ): + super().__init__() + self.block = block + + def forward(self, x, average_language_embedding): + return self.block(x) + + +def unpack_tuple(fn: Callable[[Any], Tuple[Any]]) -> Callable[[Any], Any]: + """Utility function for monkey-patching functions.""" + + def wrapper(*args: Any, **kwargs: Any) -> Any: + result = fn(*args, **kwargs) + return result[0] if isinstance(result, tuple) else result + + return wrapper + + +class FiLMedVisionTransformer(VisionTransformer): + """ + Wrapper for timm.models.vision_transformer.VisionTransformer that overrides functions to enable infusing language + embeddings into visual embeddings via FiLM. + """ + + def _intermediate_layers( + self, + x: torch.Tensor, + language_embeddings: torch.Tensor, + n: Union[int, Sequence] = 1, + ): + """ + Copy of timm.models.vision_transformer.VisionTransformer._intermediate_layers() with modifications + to take in language embeddings as additional input. + """ + outputs, num_blocks = [], len(self.blocks) + take_indices = set(range(num_blocks - n, num_blocks) if isinstance(n, int) else n) + + # forward pass + x = self.patch_embed(x) + x = self._pos_embed(x) + x = self.patch_drop(x) + x = self.norm_pre(x) + for i, blk in enumerate(self.blocks): + x = blk(x, language_embeddings) # Modified to receive language_embeddings + if i in take_indices: + outputs.append(x) + + return outputs + + def get_intermediate_layers( + self, + x: torch.Tensor, + language_embeddings: torch.Tensor, + n: Union[int, Sequence] = 1, + reshape: bool = False, + return_prefix_tokens: bool = False, + norm: bool = False, + ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]: + """ + Copy of timm.models.vision_transformer.VisionTransformer.get_intermediate_layers() with modifications + to allow language embeddings as additional input. + """ + # take last n blocks if n is an int, if in is a sequence, select by matching indices + outputs = self._intermediate_layers(x, language_embeddings, n) + if norm: + outputs = [self.norm(out) for out in outputs] + prefix_tokens = [out[:, 0 : self.num_prefix_tokens] for out in outputs] + outputs = [out[:, self.num_prefix_tokens :] for out in outputs] + + if reshape: + grid_size = self.patch_embed.grid_size + outputs = [ + out.reshape(x.shape[0], grid_size[0], grid_size[1], -1).permute(0, 3, 1, 2).contiguous() + for out in outputs + ] + + if return_prefix_tokens: + return tuple(zip(outputs, prefix_tokens)) + return tuple(outputs) + + +class FiLMedPrismaticVisionBackbone(nn.Module): + """ + Wrapper for OpenVLA's vision backbone that implements feature-wise linear modulation (FiLM). + + Wraps the Vision Transformers in the vision backbone to enable language conditioning through FiLM. + Supports processing 1-3 images using dual vision backbones (SigLIP + DINOv2). + """ + + def __init__( + self, + vision_backbone, + llm_dim: int = 4096, # 4096 for Llama-2 7B + ) -> None: + """ + Initializes FiLM wrapper. + + Args: + vision_backbone (PrismaticVisionBackbone): Base vision backbone. + llm_dim (int): Dimension of language model embeddings. + """ + super().__init__() + self.vision_backbone = vision_backbone + self.llm_dim = llm_dim + + # Wrap vision transformers + self._wrap_vit(self.vision_backbone.featurizer) # SigLIP + if self.vision_backbone.use_fused_vision_backbone: + self._wrap_vit(self.vision_backbone.fused_featurizer) # DINOv2 + + def _wrap_vit(self, vit) -> None: + """ + Creates wrapper around an individual vision transformer to allow for infusion of language inputs. + + Args: + vit (VisionTransformer): Original vision transformer. + """ + # Wrap vision transformer blocks + block_wrappers = [] + for block in vit.blocks: + block_wrappers.append( + FiLMedVisionTransformerBlock(block=block, vision_dim=vit.num_features, llm_dim=self.llm_dim) + ) + vit.blocks = nn.Sequential(*block_wrappers) + + # Wrap vision transformer with new class that overrides functions used for forward pass + vit.__class__ = FiLMedVisionTransformer + vit.forward = unpack_tuple(partial(vit.get_intermediate_layers, n={len(vit.blocks) - 2})) + + def get_num_patches(self) -> int: + """Returns the number of vision patches output by the vision backbone.""" + return self.vision_backbone.get_num_patches() + + def get_num_images_in_input(self) -> int: + """Returns the number of input images for the vision backbone.""" + return self.vision_backbone.get_num_images_in_input() + + def set_num_images_in_input(self, num_images_in_input: int) -> None: + """Sets the number of input images for the vision backbone.""" + self.vision_backbone.set_num_images_in_input(num_images_in_input) + + def forward(self, pixel_values: torch.Tensor, language_embeddings: torch.Tensor) -> torch.Tensor: + """ + Implements the forward pass for the vision backbone with FiLM to infuse language inputs into visual features. + + Identical to PrismaticVisionBackbone.forward() except that language embeddings are also used as input. + + Args: + pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W). + language_embeddings (torch.Tensor): Language embeddings for the task description, (B, seq_len, llm_dim). + """ + # For FiLM: Average the language embeddings of the task description + average_language_embedding = language_embeddings.mean(dim=1) + + if self.get_num_images_in_input() == 1: + if not self.vision_backbone.use_fused_vision_backbone: + return self.vision_backbone.featurizer(pixel_values, average_language_embedding) + + # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack + img, img_fused = torch.split(pixel_values, [3, 3], dim=1) + patches = self.vision_backbone.featurizer(img, average_language_embedding) + patches_fused = self.vision_backbone.fused_featurizer(img_fused, average_language_embedding) + + return torch.cat([patches, patches_fused], dim=2) + + else: + assert self.vision_backbone.use_fused_vision_backbone, "Multi-image inputs require using fused backbone!" + + # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2) + images = torch.split(pixel_values, [6] * self.get_num_images_in_input(), dim=1) + + # Process each image and collect patches + all_patches = [] + for img in images: + # Split each image further into two stacks of channels (each with 3 channels) + img_regular, img_fused = torch.split(img, [3, 3], dim=1) + + # Get patches from both SigLIP and DINOv2 vision transformers + patches = self.vision_backbone.featurizer(img_regular, average_language_embedding) + patches_fused = self.vision_backbone.fused_featurizer(img_fused, average_language_embedding) + + # Concatenate SigLIP and DINOv2 patches along the hidden dimension + combined_patches = torch.cat([patches, patches_fused], dim=2) + all_patches.append(combined_patches) + + # Concatenate all patches along the patch dimension + return torch.cat(all_patches, dim=1) diff --git a/prismatic/models/projectors.py b/prismatic/models/projectors.py new file mode 100644 index 000000000..ea20dade1 --- /dev/null +++ b/prismatic/models/projectors.py @@ -0,0 +1,49 @@ +"""Implementation of additional projectors for additional inputs to the VLA models.""" +import torch +import torch.nn as nn + + +class ProprioProjector(nn.Module): + """ + Projects proprio state inputs into the LLM's embedding space. + """ + def __init__(self, llm_dim: int, proprio_dim: int) -> None: + super().__init__() + self.llm_dim = llm_dim + self.proprio_dim = proprio_dim + + self.fc1 = nn.Linear(self.proprio_dim, self.llm_dim, bias=True) + self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True) + self.act_fn1 = nn.GELU() + + def forward(self, proprio: torch.Tensor = None) -> torch.Tensor: + # proprio: (bsz, proprio_dim) + projected_features = self.fc1(proprio) + projected_features = self.act_fn1(projected_features) + projected_features = self.fc2(projected_features) + return projected_features + + +class NoisyActionProjector(nn.Module): + """ + [Diffusion] Projects noisy action inputs into the LLM's embedding space. + + Note that since each action is tokenized into 7 tokens in OpenVLA (rather + than having 1 token per action), each noisy action token will have dimension 1 + instead of 7. + """ + def __init__(self, llm_dim: int) -> None: + super().__init__() + self.llm_dim = llm_dim + self.action_token_dim = 1 + + self.fc1 = nn.Linear(self.action_token_dim, self.llm_dim, bias=True) + self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True) + self.act_fn1 = nn.GELU() + + def forward(self, noisy_actions: torch.Tensor = None) -> torch.Tensor: + # noisy_actions: (bsz, num_action_tokens=chunk_len*action_dim, 1) + projected_features = self.fc1(noisy_actions) + projected_features = self.act_fn1(projected_features) + projected_features = self.fc2(projected_features) + return projected_features diff --git a/prismatic/training/strategies/base_strategy.py b/prismatic/training/strategies/base_strategy.py index 018ee41cf..ba4fc9428 100644 --- a/prismatic/training/strategies/base_strategy.py +++ b/prismatic/training/strategies/base_strategy.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import Callable, Optional +import numpy as np import torch import torch.distributed as dist from torch.utils.data import DataLoader, Dataset, DistributedSampler, IterableDataset @@ -21,11 +22,22 @@ from prismatic.models.vlms import PrismaticVLM from prismatic.overwatch import initialize_overwatch from prismatic.training.metrics import Metrics, VLAMetrics +from prismatic.training.train_utils import ( + compute_actions_l1_loss, + compute_token_accuracy, + get_current_action_mask, + get_next_actions_mask, +) from prismatic.util import check_bloat16_supported from prismatic.util.batching_utils import SplitModalitySampler from prismatic.util.data_utils import PaddedCollatorForActionPrediction, PaddedCollatorForLanguageModeling from prismatic.vla.action_tokenizer import ActionTokenizer +# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels) +from prismatic.vla.constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, NUM_ACTIONS_CHUNK, IGNORE_INDEX +NEWLINE_INDEX = 13 # '\n' +STOP_INDEX = 2 # '' + # Initialize Overwatch =>> Wraps `logging.Logger` overwatch = initialize_overwatch(__name__) @@ -300,36 +312,48 @@ def run_vla_training( metrics.commit(loss=loss) loss.backward() - # === Compute Action Token Accuracy & L1 Loss === - - # To compute action token accuracy, we need to identify the locations of the action tokens - # in both `output.logits` and `batch["labels"]`. We know that when "right" padding, we - # insert `self.vlm.vision_backbone.num_patches` at index 1. - # - # Computing `action_prediction_accuracy` is then pretty straightforward: - # 1) Extract "aligned" predictions & labels - # 2) Compute boolean "mask" where "labels > 2" (where 2 is ID for `EOS_TOKEN`) - # => If masking out EOS, then it's just "labels != -100 (IGNORE_INDEX) - # 3) Compute masked accuracy as `(preds == logits) & mask` --> sum/divide by # unmasked! - action_preds = output.logits[:, self.vlm.vision_backbone.num_patches : -1].argmax(dim=2) - action_gt = batch["labels"][:, 1:].to(action_preds.device) - mask = action_gt > action_tokenizer.action_token_begin_idx + # Get predicted and ground-truth token IDs + predicted_token_ids = output.logits[:, self.vlm.vision_backbone.num_patches : -1].argmax(dim=2) + ground_truth_token_ids = batch["labels"][:, 1:].to(predicted_token_ids.device) + + ####################################################################### + # === Compute Current Action Token Accuracy & L1 Loss === + ####################################################################### + + # Get current action mask: Target the first ACTION_DIM non-ignore tokens + current_action_mask = get_current_action_mask(ground_truth_token_ids) # Compute Accuracy - correct_preds = (action_preds == action_gt) & mask - action_accuracy = correct_preds.sum().float() / mask.sum().float() + action_accuracy = compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask=current_action_mask) # Compute L1 Loss on Predicted (Continuous) Actions - continuous_actions_pred = torch.tensor( - action_tokenizer.decode_token_ids_to_actions(action_preds[mask].cpu().numpy()) - ) - continuous_actions_gt = torch.tensor( - action_tokenizer.decode_token_ids_to_actions(action_gt[mask].cpu().numpy()) - ) - action_l1_loss = torch.nn.functional.l1_loss(continuous_actions_pred, continuous_actions_gt) + action_l1_loss = compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask=current_action_mask) + + ####################################################################### + # === Compute Next Actions Token Accuracy & L1 Loss === + ####################################################################### + + # Get next actions mask: Target all tokens after the first ACTION_DIM non-ignore tokens (excluding the last token, which is the stop token) + next_actions_mask = get_next_actions_mask(ground_truth_token_ids) + + # Compute Accuracy + next_actions_accuracy = compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask=next_actions_mask) + + # Compute L1 Loss on Predicted (Continuous) Actions + next_actions_l1_loss = compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask=next_actions_mask) + + ####################################################################### + # === Log === + ####################################################################### # Commit Metrics - metrics.commit(action_accuracy=action_accuracy, l1_loss=action_l1_loss, update_step_time=True) + metrics.commit( + action_accuracy=action_accuracy, + l1_loss=action_l1_loss, + next_actions_accuracy=next_actions_accuracy, + next_actions_l1_loss=next_actions_l1_loss, + update_step_time=True, + ) # Compute metrics per dataset --> only on rank_zero since we don't log them on other workers anyways if overwatch.is_rank_zero(): @@ -338,21 +362,25 @@ def run_vla_training( for ds in datasets: ds_mask = torch.tensor([elem == ds for elem in batch["dataset_names"]]) action_accuracy_ds = correct_preds[ds_mask].sum().float() / mask[ds_mask].sum().float() - continuous_actions_pred_ds = torch.tensor( + pred_continuous_actions_ds = torch.tensor( action_tokenizer.decode_token_ids_to_actions( - action_preds[ds_mask][mask[ds_mask]].cpu().numpy() + predicted_token_ids[ds_mask][mask[ds_mask]].cpu().numpy() ) ) continuous_actions_gt_ds = torch.tensor( action_tokenizer.decode_token_ids_to_actions( - action_gt[ds_mask][mask[ds_mask]].cpu().numpy() + ground_truth_token_ids[ds_mask][mask[ds_mask]].cpu().numpy() ) ) action_l1_loss_ds = torch.nn.functional.l1_loss( - continuous_actions_pred_ds, continuous_actions_gt_ds + pred_continuous_actions_ds, continuous_actions_gt_ds ) metrics.commit_for_dataset( - dataset_name=ds.decode(), action_accuracy=action_accuracy_ds, l1_loss=action_l1_loss_ds + dataset_name=ds.decode(), + action_accuracy=action_accuracy_ds, + l1_loss=action_l1_loss_ds, + next_actions_accuracy=next_actions_accuracy, + next_actions_l1_loss=next_actions_l1_loss, ) # === Gradient Step === diff --git a/prismatic/training/train_utils.py b/prismatic/training/train_utils.py new file mode 100644 index 000000000..0c546885d --- /dev/null +++ b/prismatic/training/train_utils.py @@ -0,0 +1,56 @@ +"""Utils for training/fine-tuning scripts.""" + +import torch + +from prismatic.vla.constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX + + +def get_current_action_mask(token_ids): + # Create a tensor marking positions of IGNORE_INDEX + newline_positions = token_ids != IGNORE_INDEX + + # Calculate cumulative sum to identify regions between newlines + cumsum = torch.cumsum(newline_positions, dim=1) + + # Create the mask + mask = (1 <= cumsum) & (cumsum <= ACTION_DIM) + + # Extract the action part only + action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX + mask = action_tokens_only_mask * mask + + return mask + + +def get_next_actions_mask(token_ids): + # Create a tensor marking positions of IGNORE_INDEX + newline_positions = token_ids != IGNORE_INDEX + + # Calculate cumulative sum to identify regions between newlines + cumsum = torch.cumsum(newline_positions, dim=1) + + # Create the mask + mask = cumsum > ACTION_DIM + + # Extract the action part only + action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX + mask = action_tokens_only_mask * mask + + return mask + + +def compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask): + correct_preds = (predicted_token_ids == ground_truth_token_ids) & mask + accuracy = correct_preds.sum().float() / mask.sum().float() + return accuracy + + +def compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask): + pred_continuous_actions = torch.tensor( + action_tokenizer.decode_token_ids_to_actions(predicted_token_ids[mask].cpu().numpy()) + ) + true_continuous_actions = torch.tensor( + action_tokenizer.decode_token_ids_to_actions(ground_truth_token_ids[mask].cpu().numpy()) + ) + l1_loss = torch.nn.functional.l1_loss(pred_continuous_actions, true_continuous_actions) + return l1_loss diff --git a/prismatic/util/data_utils.py b/prismatic/util/data_utils.py index cbed9603e..83fd82f19 100644 --- a/prismatic/util/data_utils.py +++ b/prismatic/util/data_utils.py @@ -5,11 +5,20 @@ """ from dataclasses import dataclass -from typing import Callable, Dict, Sequence, Tuple - +from typing import ( + Any, + Callable, + Dict, + Sequence, + Tuple, +) + +import numpy as np import torch +from PIL import Image from torch.nn.utils.rnn import pad_sequence + # HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels) IGNORE_INDEX = -100 @@ -26,6 +35,33 @@ def tree_map_with_key(fn: Callable, tree: dict, keys: Sequence = ()) -> dict: } +def greyscale_float_tensor_preprocessing_wrapper( + transform_fn: Callable[[Image.Image | np.ndarray], Any] +) -> Callable[[torch.Tensor | Image.Image | np.ndarray], Any]: + """ + Wraps a transform function that expects PIL Images to work with greyscale float tensors. + + Args: + transform_fn: A function that takes a PIL Image / numpy arrayand transforms it + + Returns: + A function that can handle greyscale float tensors, PIL Images, and numpy arrays + """ + def wrapper(tensor_image): + if isinstance(tensor_image, torch.Tensor): + # Convert tensor to PIL Image + # The tensor is expected to be [C, H, W] with values in [0, 1] + img_array = (tensor_image.permute(1, 2, 0).numpy() * 255).astype(np.uint8) + pil_image = Image.fromarray(img_array) + # Apply the transform + return transform_fn(pil_image) + else: + # If it's something else, let the transform function handle it + return transform_fn(tensor_image) + + return wrapper + + @dataclass class PaddedCollatorForLanguageModeling: model_max_length: int @@ -123,20 +159,33 @@ def __call__(self, instances: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, to # Stack all `pixel_values` --> depending on type is torch.Tensor or Dict[str, torch.Tensor] if isinstance(pixel_values[0], torch.Tensor): - pixel_values = torch.stack(pixel_values) - elif isinstance(pixel_values[0], dict): - pixel_values = { - k: torch.stack([pixel_values[idx][k] for idx in range(len(input_ids))]) for k in pixel_values[0] - } + if "pixel_values_wrist" in instances[0]: + pixel_values_wrist = [instance["pixel_values_wrist"] for instance in instances] + pixel_values = torch.cat((torch.stack(pixel_values), torch.stack(pixel_values_wrist)), dim=1) + else: + pixel_values = torch.stack(pixel_values) else: raise ValueError(f"Unsupported `pixel_values` type = {type(pixel_values)}") + # Stack all actions + actions = [torch.from_numpy(np.copy(instance["actions"])) for instance in instances] + actions = torch.stack(actions) + + # Stack proprio + if "proprio" in instances[0]: + proprio = [instance["proprio"] for instance in instances] + proprio = torch.Tensor(np.squeeze(np.stack(proprio))) + else: + proprio = None + output = dict( pixel_values=pixel_values, + proprio=proprio, input_ids=input_ids, attention_mask=attention_mask, labels=labels, + actions=actions, ) if dataset_names is not None: output["dataset_names"] = dataset_names - return output + return output \ No newline at end of file diff --git a/prismatic/util/extern/__init__.py b/prismatic/util/extern/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/prismatic/util/extern/hf/__init__.py b/prismatic/util/extern/hf/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/prismatic/util/extern/hf/lerobot_utils.py b/prismatic/util/extern/hf/lerobot_utils.py new file mode 100644 index 000000000..e1434bdac --- /dev/null +++ b/prismatic/util/extern/hf/lerobot_utils.py @@ -0,0 +1,314 @@ +""" +lerobot_utils.py + +Utilities for working with LeRobotDataset v2.0. +""" + +from dataclasses import dataclass +from typing import ( + Optional, + Sequence, + Type, +) + +import numpy as np +import torch +from torch.nn.utils.rnn import pad_sequence +from transformers import PreTrainedTokenizerBase + +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset + +from prismatic.vla.constants import ( + ACTION_DIM, + ACTION_PROPRIO_NORMALIZATION_TYPE, + IGNORE_INDEX, +) +from prismatic.vla.action_tokenizer import ActionTokenizer +from prismatic.models.backbones.llm.prompting import PurePromptBuilder + + +def create_action_norm_stats_dict_from_lerobot_dataset( + dataset: LeRobotDataset, +) -> dict[str, dict[str, list[float]]]: + """ + Get statistics for unnormalizing actions from a v2.0 LeRobotDataset. + """ + if ACTION_PROPRIO_NORMALIZATION_TYPE != "bounds_q99": + raise NotImplementedError( + "For now, only q01/q99 normalization is supported " + "for OpenVLA-OFT with LeRobotDataset v2.0" + ) + assert ( + "action" in dataset.meta.stats + and "q01" in dataset.meta.stats["action"] + and "q99" in dataset.meta.stats["action"] + and len(dataset.meta.stats["action"]["q01"]) == ACTION_DIM + and len(dataset.meta.stats["action"]["q99"]) == ACTION_DIM + ), "Dataset must have q01 and q99 stored for each action dimension" + + action_norm_stats = { + "q01": dataset.meta.stats["action"]["q01"].tolist(), + "q99": dataset.meta.stats["action"]["q99"].tolist(), + } + return action_norm_stats + + +def create_rlds_dataset_stats_dict_from_lerobot_dataset( + dataset: LeRobotDataset, + dataset_name: str, +) -> dict[str, dict[str, float | list[float] | dict]]: + """ + Create a dictionary of statistics from a v2.0 LeRobotDataset that stores + action normalization statistics. + """ + + try: + action_norm_stats = \ + create_action_norm_stats_dict_from_lerobot_dataset(dataset) + except Exception as e: + raise ValueError( + f"Couldn't retrieve action norm stats from dataset: {e}" + ) from e + + dataset_stats = { + dataset_name: { + # Copy all action statistics + "action": action_norm_stats, + # Add trajectory/transition counts + "num_trajectories": dataset.num_episodes, + "num_transitions": dataset.num_frames + } + } + + # Add proprioceptive statistics if available + if "proprio" in dataset.meta.stats: + dataset_stats[dataset_name]["proprio"] \ + = dataset.meta.stats["proprio"] + + # Add any other available statistics + for key, value in dataset.meta.stats.items(): + if key not in ["action", "proprio"]: + dataset_stats[dataset_name][key] = value + + return dataset_stats + + +def create_train_val_split_from_lerobot_dataset( + dataset: LeRobotDataset, + split: float = 0.1, +) -> tuple[LeRobotDataset, LeRobotDataset]: + """ + Create a trajectory-based train/val split from a LeRobotDataset. + """ + + episode_indices = list(range(dataset.num_episodes)) + np.random.shuffle(episode_indices) + split = int(np.floor(split * len(episode_indices))) + step_indices_by_episode = [ + np.arange( + start=dataset.episode_data_index['from'][ep_idx], + stop=dataset.episode_data_index['to'][ep_idx], + ) + for ep_idx in episode_indices + ] + train_indices = [ + int(idx) + for idx in np.concatenate(step_indices_by_episode[split:]) + ] + val_indices = [ + int(idx) + for idx in np.concatenate(step_indices_by_episode[:split]) + ] + + train_subset = Subset(dataset, train_indices) + val_subset = Subset(dataset, val_indices) + return train_subset, val_subset + + +@dataclass +class VLACollatorForLeRobotDataset: + """ + Collator for LeRobotDataset instances specifically for VLA training. + + This collator handles: + 1. Action tokenization + 2. Prompt construction + 3. Input/label tokenization and masking + 4. Proper batching with padding + """ + action_tokenizer: ActionTokenizer + base_tokenizer: PreTrainedTokenizerBase + prompt_builder_fn: Type[PurePromptBuilder] + pad_token_id: int + model_max_length: int + predict_stop_token: bool = True + use_wrist_image: bool = False + use_proprio: bool = False + action_norm_stats: Optional[dict[str, np.ndarray]] = None + + def __call__(self, instances: Sequence[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: + """Process and collate a batch of instances from LeRobotDataset.""" + batch_size = len(instances) + + # 1. Extract data from instances + processed_items = [] + for item in instances: + # Extract task/instruction + task = item.get("task", "") + + # === MODIFIED: Extract action_chunk === + # TODO(alan): Remove this once we have a way to get the action chunks from the dataset + # Retrieve the pre-computed action chunk (which is already a tensor from the Dataset subclass) + action_chunk_tensor = item.get("action_chunk") + if action_chunk_tensor is None: + raise NotImplementedError( + "Item missing 'action_chunk'. This collator requires the dataset to be pre-processed " + "by a script (e.g., create_dataset_with_action_chunking.py) to add this column. " + "Standard LeRobotDataset does not provide chunks directly in this format." + ) + + # Ensure the tensor has the correct dtype (float32) + action_chunk_tensor = action_chunk_tensor.to(dtype=torch.float32) + + # === Original Sanity Check (can keep) === + if action_chunk_tensor.shape[1] != ACTION_DIM: + raise ValueError(f"Action chunk dimension {action_chunk_tensor.shape[1]} does not match ACTION_DIM {ACTION_DIM}") + # ======================================== + + # === MODIFIED: Normalize the whole chunk === + # Normalize actions if stats are provided + if self.action_norm_stats is not None: + q01 = self.action_norm_stats.get("q01") + q99 = self.action_norm_stats.get("q99") + # Raise error if normalization stats are expected but incomplete + if q01 is None or q99 is None: + raise ValueError( + "'action_norm_stats' was provided, but missing " + "'q01' or 'q99' keys. Cannot normalize actions." + ) + + # Proceed with normalization only if stats are valid + q01 = torch.tensor(q01, dtype=action_chunk_tensor.dtype) + q99 = torch.tensor(q99, dtype=action_chunk_tensor.dtype) + # Apply normalization across the whole chunk tensor + normalized_action_chunk \ + = (2 * action_chunk_tensor - q01 - q99) / (q99 - q01) + else: + # Keep actions unnormalized if no stats were provided at all + raise ValueError( + "Action normalization stats (`action_norm_stats`) " + "were not provided to the collator, but normalization " + "is expected." + ) + # ========================================= + + # === MODIFIED: Tokenize the flattened chunk === + # Tokenize the flattened action chunk sequence + action_tokens = self.action_tokenizer(normalized_action_chunk.view(-1)) + # ============================================ + + # 2. Build prompt + prompt_builder = self.prompt_builder_fn("openvla") + conversation = [ + {"from": "human", "value": f"What action should the robot take to {task}?"}, + {"from": "gpt", "value": action_tokens}, + ] + for turn in conversation: + prompt_builder.add_turn(turn["from"], turn["value"]) + + # 3. Tokenize + tokenized = self.base_tokenizer( + prompt_builder.get_prompt(), + add_special_tokens=True, + return_tensors="pt" + ) + input_ids = tokenized.input_ids.squeeze(0) + + # 4. Create labels (copy input_ids) + labels = input_ids.clone() + + # 5. Mask labels (only keep action tokens for loss) + action_tokens_len = len(action_tokens) # Length based on action chunk + labels[:-action_tokens_len-1] = IGNORE_INDEX + if not self.predict_stop_token: + labels[-1] = IGNORE_INDEX + + # 6. Add to processed items + processed_item = { + "input_ids": input_ids, + "labels": labels, + "pixel_values": item.get("pixel_values") if "pixel_values" in item else item.get(next(k for k in item if "image" in k.lower())), + "actions": normalized_action_chunk # Store the potentially normalized action chunk tensor + } + + # Add dataset name if available + if "dataset_name" in item: + processed_item["dataset_name"] = item["dataset_name"] + + # Add wrist camera if used + if self.use_wrist_image and any("wrist" in k.lower() for k in item): + wrist_keys = [k for k in item if "wrist" in k.lower()] + if wrist_keys: + processed_item["pixel_values_wrist"] = torch.stack([item[k] for k in wrist_keys]) + + # Add proprioceptive data if used + if self.use_proprio and "proprio" in item: + processed_item["proprio"] = item["proprio"] + + processed_items.append(processed_item) + + # 7. Collate inputs with padding + input_ids = [item["input_ids"] for item in processed_items] + labels = [item["labels"] for item in processed_items] + + # Pad sequences + input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.pad_token_id) + labels = pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX) + + # Truncate if necessary + input_ids = input_ids[:, :self.model_max_length] + labels = labels[:, :self.model_max_length] + + # Create attention mask based on padding + attention_mask = input_ids.ne(self.pad_token_id) + + # 8. Collate images and actions + # Stack main images + pixel_values = torch.stack([item["pixel_values"] for item in processed_items]) + + # If wrist images are available, combine them with the main images + if self.use_wrist_image and all("pixel_values_wrist" in item for item in processed_items): + pixel_values_wrist = torch.stack([item["pixel_values_wrist"] for item in processed_items]) + + # Reshape wrist images if needed (from [B, num_wrist, C, H, W] to [B, num_wrist*C, H, W]) + if pixel_values_wrist.dim() == 5: # [B, num_wrist, C, H, W] + B, num_wrist, C, H, W = pixel_values_wrist.shape + pixel_values_wrist = pixel_values_wrist.view(B, num_wrist * C, H, W) + + # Concatenate main and wrist images along the channel dimension + pixel_values = torch.cat([pixel_values, pixel_values_wrist], dim=1) + + actions = torch.stack([item["actions"] for item in processed_items]) + + # 9. Build final batch + batch = { + "pixel_values": pixel_values, + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "actions": actions, + } + + # Add wrist images if available + if self.use_wrist_image and all("pixel_values_wrist" in item for item in processed_items): + batch["pixel_values_wrist"] = torch.cat([item["pixel_values_wrist"] for item in processed_items], dim=1) + + # Add proprioceptive data if available + if self.use_proprio and all("proprio" in item for item in processed_items): + batch["proprio"] = torch.stack([item["proprio"] for item in processed_items]) + + # Add dataset names if available + if all("dataset_name" in item for item in processed_items): + batch["dataset_names"] = [item["dataset_name"] for item in processed_items] + + return batch \ No newline at end of file diff --git a/prismatic/vla/constants.py b/prismatic/vla/constants.py new file mode 100644 index 000000000..73b174144 --- /dev/null +++ b/prismatic/vla/constants.py @@ -0,0 +1,96 @@ +""" +Important constants for VLA training and evaluation. + +Attempts to automatically identify the correct constants to set based on the Python command used to launch +training or evaluation. If it is unclear, defaults to using the LIBERO simulation benchmark constants. +""" +import sys +from enum import Enum + +# Llama 2 token constants +IGNORE_INDEX = -100 +ACTION_TOKEN_BEGIN_IDX = 31743 +STOP_INDEX = 2 # '' + + +# Defines supported normalization schemes for action and proprioceptive state. +class NormalizationType(str, Enum): + # fmt: off + NORMAL = "normal" # Normalize to Mean = 0, Stdev = 1 + BOUNDS = "bounds" # Normalize to Interval = [-1, 1] + BOUNDS_Q99 = "bounds_q99" # Normalize [quantile_01, ..., quantile_99] --> [-1, ..., 1] + # fmt: on + + +# Define constants for each robot platform +LIBERO_CONSTANTS = { + "NUM_ACTIONS_CHUNK": 8, + "ACTION_DIM": 7, + "PROPRIO_DIM": 8, + "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS_Q99, +} + +ALOHA_CONSTANTS = { + "NUM_ACTIONS_CHUNK": 25, + "ACTION_DIM": 14, + "PROPRIO_DIM": 14, + "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS, +} + +BRIDGE_CONSTANTS = { + "NUM_ACTIONS_CHUNK": 5, + "ACTION_DIM": 7, + "PROPRIO_DIM": 7, + "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS_Q99, +} + +UR5E_CONSTANTS = { + "NUM_ACTIONS_CHUNK": 8, + "ACTION_DIM": 7, + "PROPRIO_DIM": 0, + "ACTION_PROPRIO_NORMALIZATION_TYPE": NormalizationType.BOUNDS_Q99, +} + +# Function to detect robot platform from command line arguments +def detect_robot_platform(): + cmd_args = " ".join(sys.argv).lower() + + if "libero" in cmd_args: + return "LIBERO" + elif "aloha" in cmd_args: + return "ALOHA" + elif "bridge" in cmd_args: + return "BRIDGE" + elif "ur5e" in cmd_args: + return "UR5E" + else: + # Default to LIBERO if unclear + return "LIBERO" + + +# Determine which robot platform to use +ROBOT_PLATFORM = detect_robot_platform() + +# Set the appropriate constants based on the detected platform +if ROBOT_PLATFORM == "LIBERO": + constants = LIBERO_CONSTANTS +elif ROBOT_PLATFORM == "ALOHA": + constants = ALOHA_CONSTANTS +elif ROBOT_PLATFORM == "BRIDGE": + constants = BRIDGE_CONSTANTS +elif ROBOT_PLATFORM == "UR5E": + constants = UR5E_CONSTANTS + +# Assign constants to global variables +NUM_ACTIONS_CHUNK = constants["NUM_ACTIONS_CHUNK"] +ACTION_DIM = constants["ACTION_DIM"] +PROPRIO_DIM = constants["PROPRIO_DIM"] +ACTION_PROPRIO_NORMALIZATION_TYPE = constants["ACTION_PROPRIO_NORMALIZATION_TYPE"] + +# Print which robot platform constants are being used (for debugging) +print(f"Using {ROBOT_PLATFORM} constants:") +print(f" NUM_ACTIONS_CHUNK = {NUM_ACTIONS_CHUNK}") +print(f" ACTION_DIM = {ACTION_DIM}") +print(f" PROPRIO_DIM = {PROPRIO_DIM}") +print(f" ACTION_PROPRIO_NORMALIZATION_TYPE = {ACTION_PROPRIO_NORMALIZATION_TYPE}") +print("If needed, manually set the correct constants in `prismatic/vla/constants.py`!") diff --git a/prismatic/vla/datasets/datasets.py b/prismatic/vla/datasets/datasets.py index 539b4144d..34d2de93e 100644 --- a/prismatic/vla/datasets/datasets.py +++ b/prismatic/vla/datasets/datasets.py @@ -19,13 +19,9 @@ from prismatic.models.backbones.vision import ImageTransform from prismatic.util.data_utils import tree_map from prismatic.vla.action_tokenizer import ActionTokenizer +from prismatic.vla.constants import ACTION_DIM, ACTION_PROPRIO_NORMALIZATION_TYPE, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX, NUM_ACTIONS_CHUNK, PROPRIO_DIM, STOP_INDEX from prismatic.vla.datasets.rlds import make_interleaved_dataset, make_single_dataset from prismatic.vla.datasets.rlds.oxe import OXE_NAMED_MIXTURES, get_oxe_dataset_kwargs_and_weights -from prismatic.vla.datasets.rlds.utils.data_utils import NormalizationType - -# HuggingFace Default / LLaMa-2 IGNORE_INDEX (for labels) -IGNORE_INDEX = -100 - @dataclass class RLDSBatchTransform: @@ -34,18 +30,31 @@ class RLDSBatchTransform: image_transform: ImageTransform prompt_builder_fn: Type[PromptBuilder] predict_stop_token: bool = True + use_wrist_image: bool = False + use_proprio: bool = False def __call__(self, rlds_batch: Dict[str, Any]) -> Dict[str, Any]: """Converts a RLDS batch to the format expected by the OpenVLA collator/models.""" - dataset_name, action = rlds_batch["dataset_name"], rlds_batch["action"][0] + dataset_name, current_action = rlds_batch["dataset_name"], rlds_batch["action"][0] img = Image.fromarray(rlds_batch["observation"]["image_primary"][0]) lang = rlds_batch["task"]["language_instruction"].decode().lower() + actions = rlds_batch["action"] # Construct Chat-based Prompt =>> Input is default query + language instruction, output are the action tokens prompt_builder = self.prompt_builder_fn("openvla") + + # Get future action chunk + future_actions = rlds_batch["action"][1:] + future_actions_string = ''.join(self.action_tokenizer(future_actions)) + + # Get action chunk string + current_action_string = self.action_tokenizer(current_action) + action_chunk_string = current_action_string + future_actions_string + action_chunk_len = len(action_chunk_string) + conversation = [ {"from": "human", "value": f"What action should the robot take to {lang}?"}, - {"from": "gpt", "value": self.action_tokenizer(action)}, + {"from": "gpt", "value": action_chunk_string}, ] for turn in conversation: prompt_builder.add_turn(turn["from"], turn["value"]) @@ -60,11 +69,26 @@ def __call__(self, rlds_batch: Dict[str, Any]) -> Dict[str, Any]: pixel_values = self.image_transform(img) # [CRITICAL] We do not want to take the loss for anything but the predicted action tokens! - labels[: -(len(action) + 1)] = IGNORE_INDEX + labels[: -(action_chunk_len + 1)] = IGNORE_INDEX if not self.predict_stop_token: labels[-1] = IGNORE_INDEX - return dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels, dataset_name=dataset_name) + return_dict = dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels, dataset_name=dataset_name, actions=actions) + + # Add additional inputs + if self.use_wrist_image: + all_wrist_pixels = [] + for k in rlds_batch["observation"].keys(): + if "wrist" in k: + img_wrist = Image.fromarray(rlds_batch["observation"][k][0]) + pixel_values_wrist = self.image_transform(img_wrist) + all_wrist_pixels.append(pixel_values_wrist) + return_dict["pixel_values_wrist"] = torch.cat(all_wrist_pixels, dim=0) + if self.use_proprio and "proprio" in rlds_batch["observation"]: + proprio = rlds_batch["observation"]["proprio"] + return_dict["proprio"] = proprio + + return return_dict class RLDSDataset(IterableDataset): @@ -89,19 +113,24 @@ def __init__( mixture_spec = [(self.data_mix, 1.0)] # fmt: off + if "aloha" in self.data_mix: + load_camera_views = ("primary", "left_wrist", "right_wrist") + else: + load_camera_views = ("primary", "wrist") + per_dataset_kwargs, weights = get_oxe_dataset_kwargs_and_weights( self.data_root_dir, mixture_spec, - load_camera_views=("primary",), + load_camera_views=load_camera_views, load_depth=False, - load_proprio=False, + load_proprio=True, load_language=True, - action_proprio_normalization_type=NormalizationType.BOUNDS_Q99, + action_proprio_normalization_type=ACTION_PROPRIO_NORMALIZATION_TYPE, ) rlds_config = dict( traj_transform_kwargs=dict( window_size=1, # If we wanted to feed / predict more than one step - future_action_window_size=0, # For action chunking + future_action_window_size=NUM_ACTIONS_CHUNK-1, # For action chunking skip_unlabeled=True, # Skip trajectories without language labels goal_relabeling_strategy="uniform", # Goals are currently unused ), @@ -176,6 +205,7 @@ def __iter__(self) -> Dict[str, Any]: ] yield out +### class DummyDataset(Dataset): def __init__( @@ -203,17 +233,28 @@ def __len__(self): return 10000 def __getitem__(self, idx): - # TODO =>> Load image, action and instruction from disk -- we use dummy values - image = Image.fromarray(np.asarray(np.random.rand(224, 224, 3) * 255.0, dtype=np.uint8)) + """Get a single training example.""" + # Generate random image, action and instruction + image = Image.fromarray( + np.asarray(np.random.rand(224, 224, 3) * 255.0, dtype=np.uint8) + ) action = np.asarray(np.random.rand(7), dtype=np.float32) instruction = "do something spectacular" - # Add instruction to VLA prompt + # Build conversation prompt prompt_builder = self.prompt_builder_fn("openvla") conversation = [ - {"from": "human", "value": f"What action should the robot take to {instruction}?"}, - {"from": "gpt", "value": self.action_tokenizer(action)}, + { + "from": "human", + "value": f"What action should the robot take to {instruction}?" + }, + { + "from": "gpt", + "value": self.action_tokenizer(action) + } ] + + # Add conversation turns to prompt builder for turn in conversation: prompt_builder.add_turn(turn["from"], turn["value"]) @@ -230,3 +271,132 @@ def __getitem__(self, idx): labels[: -(len(action) + 1)] = IGNORE_INDEX return dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels) + + +from typing import Callable + +from lerobot.common.datasets.lerobot_dataset import ( + LeRobotDataset, + LeRobotDatasetMetadata, +) + + +class RLDSLeRobotDataset(LeRobotDataset): + + def __init__( + self, + repo_id: str, + action_tokenizer: ActionTokenizer, + base_tokenizer: PreTrainedTokenizerBase, + image_transform: ImageTransform, + prompt_builder_fn: Type[PromptBuilder], + *, + root: str | Path | None = None, + episodes: list[int] | None = None, + image_transforms: Callable | None = None, + delta_timestamps: dict[list[float]] | None = None, + tolerance_s: float = 1e-4, + download_videos: bool = True, + local_files_only: bool = False, + video_backend: str | None = None, + ) -> None: + super().__init__( + repo_id, + root, + episodes, + image_transforms, + delta_timestamps, + tolerance_s, + download_videos, + local_files_only, + video_backend, + ) + assert isinstance(self.meta, LeRobotDatasetMetadata) + + self.action_tokenizer = action_tokenizer + self.base_tokenizer = base_tokenizer + self.image_transform = image_transform + self.prompt_builder_fn = prompt_builder_fn + + # NOTE: We expect the dataset to store statistics for action de-normalization: + # 1/100st quantile of each action under "q01" and 99/100th quantile under "q99". + print(self.meta.stats) + self.dataset_statistics = { + "rlds_lerobot_dataset": { + "action": { + "q01": np.array(self.meta.stats["action"]["q01"]), + "q99": np.array(self.meta.stats["action"]["q99"]), + } + } + } + + # NOTE: This is hardcoded as a social contract. + # This is the only key to image observations that will be used. + self.obs_image_key = "observation.images.side" + + + def __len__(self): + return self.num_frames + + # Retrieves a single (instruction, image, action) triple from the dataset. + def __getitem__(self, idx): + + hf_item = super().__getitem__(idx) + + # Retrieve image observation. + img_array: np.ndarray = ( + hf_item[self.obs_image_key] + .permute(1, 2, 0) + .numpy() * 255 + ).astype(np.uint8) + image = Image.fromarray(img_array) + + # Retrieve instruction. + task_idx: torch.Tensor = hf_item["task_index"] + task_idx: int = task_idx.item() + instruction = self.meta.tasks[task_idx] + + # Retrieve action. + action: torch.Tensor = torch.cat([ + hf_item["action.pose"], + hf_item["action.gripper"].unsqueeze(0) + ]) + + qs = self.dataset_statistics["rlds_lerobot_dataset"]["action"] + q01, q99 = np.array(qs["q01"]), np.array(qs["q99"]) + action = (2*action - q01 - q99) / (q99 - q01) # normalize to [-1, 1] + action: str = self.action_tokenizer(action) + + # Add instruction to VLA prompt. + prompt_builder = self.prompt_builder_fn("openvla") + conversation = [ + { + "from": "human", + "value": f"Hey I need the robot to do this: {instruction}. We upgraded from the old pincer gripper to this new suction cup end effector - it's that blue circular cup with the yellow ring at the end of the silver arm. Big difference is we can't tell if it's got a good seal just by looking at it (unlike before where we could see the gripper fingers close). Also the suction cup needs a flat surface to grip well, and we need enough vacuum pressure for different weights. Sometimes we need to wiggle it a bit to break the seal when releasing too. What's the best way to handle this with the new setup?" + }, + { + "from": "gpt", + "value": f"{action}" + }, + ] + for turn in conversation: + prompt_builder.add_turn(turn["from"], turn["value"]) + prompt = prompt_builder.get_prompt() + + # Tokenize (w/ `base_tokenizer`) + input_ids = self.base_tokenizer( + prompt, + add_special_tokens=True + ).input_ids + labels = list(input_ids) + + # Tensorize =>> Run Image Transform to get `pixel_values` =>> Return + # =>> IMPORTANT :: IF WE'RE USING HF .forward(..., labels=labels), SHIFTING HAPPENS _INSIDE_ MODEL! + input_ids, labels = torch.tensor(input_ids), torch.tensor(labels) + pixel_values = self.image_transform(image) + + # [CRITICAL] We do not want to take the loss for anything but the predicted action tokens! + labels[: -(len(action) + 1)] = IGNORE_INDEX + + return dict(pixel_values=pixel_values, input_ids=input_ids, labels=labels) + \ No newline at end of file diff --git a/prismatic/vla/datasets/rlds/dataset.py b/prismatic/vla/datasets/rlds/dataset.py index e9bcd93eb..f07215a2d 100644 --- a/prismatic/vla/datasets/rlds/dataset.py +++ b/prismatic/vla/datasets/rlds/dataset.py @@ -16,10 +16,10 @@ import tensorflow_datasets as tfds from prismatic.overwatch import initialize_overwatch +from prismatic.vla.constants import ACTION_DIM, ACTION_PROPRIO_NORMALIZATION_TYPE, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX, NUM_ACTIONS_CHUNK, PROPRIO_DIM, STOP_INDEX from prismatic.vla.datasets.rlds import obs_transforms, traj_transforms from prismatic.vla.datasets.rlds.utils import goal_relabeling, task_augmentation from prismatic.vla.datasets.rlds.utils.data_utils import ( - NormalizationType, allocate_threads, get_dataset_statistics, normalize_action_and_proprio, @@ -47,7 +47,7 @@ def make_dataset_from_rlds( depth_obs_keys: Dict[str, Optional[str]] = {}, state_obs_keys: List[Optional[str]] = (), language_key: Optional[str] = None, - action_proprio_normalization_type: NormalizationType = NormalizationType.NORMAL, + action_proprio_normalization_type: ACTION_PROPRIO_NORMALIZATION_TYPE, dataset_statistics: Optional[Union[dict, str]] = None, absolute_action_mask: Optional[List[bool]] = None, action_normalization_mask: Optional[List[bool]] = None, @@ -231,10 +231,7 @@ def restructure(traj): dataset_statistics["action"]["mask"] = np.array(action_normalization_mask) # construct the dataset - if "val" not in builder.info.splits: - split = "train[:95%]" if train else "train[95%:]" - else: - split = "train" if train else "val" + split = "train" if train else "val" dataset = dl.DLataset.from_rlds(builder, split=split, shuffle=shuffle, num_parallel_reads=num_parallel_reads) diff --git a/prismatic/vla/datasets/rlds/oxe/configs.py b/prismatic/vla/datasets/rlds/oxe/configs.py index 2b8dcb931..3222e023b 100644 --- a/prismatic/vla/datasets/rlds/oxe/configs.py +++ b/prismatic/vla/datasets/rlds/oxe/configs.py @@ -72,21 +72,21 @@ class ActionEncoding(IntEnum): "bridge_oxe": { # Version of Bridge V2 in Open X-Embodiment mixture "image_obs_keys": {"primary": "image", "secondary": "image_1", "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "bridge_orig": { # Original version of Bridge V2 from project website "image_obs_keys": {"primary": "image_0", "secondary": "image_1", "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "bridge_dataset": { # Original version of Bridge V2 from project website "image_obs_keys": {"primary": "image_0", "secondary": "image_1", "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -198,7 +198,7 @@ class ActionEncoding(IntEnum): "nyu_rot_dataset_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -209,7 +209,7 @@ class ActionEncoding(IntEnum): "wrist": "wrist_image", }, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -286,7 +286,7 @@ class ActionEncoding(IntEnum): "ucsd_pick_and_place_dataset_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -327,14 +327,14 @@ class ActionEncoding(IntEnum): "utokyo_pr2_opening_fridge_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "utokyo_pr2_tabletop_manipulation_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -359,7 +359,7 @@ class ActionEncoding(IntEnum): "robo_net": { "image_obs_keys": {"primary": "image", "secondary": "image1", "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -387,14 +387,14 @@ class ActionEncoding(IntEnum): "stanford_mask_vit_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "tokyo_u_lsmo_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -422,14 +422,14 @@ class ActionEncoding(IntEnum): "asu_table_top_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "stanford_robocook_converted_externally_to_rlds": { "image_obs_keys": {"primary": "image_1", "secondary": "image_2", "wrist": None}, "depth_obs_keys": {"primary": "depth_1", "secondary": "depth_2", "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -505,7 +505,7 @@ class ActionEncoding(IntEnum): "cmu_stretch": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["eef_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -594,42 +594,42 @@ class ActionEncoding(IntEnum): "tdroid_carrot_in_bowl": { # "put carrot in bowl" task, 50 demos @ 5 Hz control "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "tdroid_pour_corn_in_pot": { # "pour corn from red bowl into steel pot" task, 50 demos @ 5 Hz control "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "tdroid_flip_pot_upright": { # "flip pot upright" task, 10 demos @ 5 Hz control "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "tdroid_move_object_onto_plate": { # "move onto plate" task, 150 demos @ 5 Hz control "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "tdroid_knock_object_over": { # "knock over" task, 70 demos @ 5 Hz control "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "tdroid_cover_object_with_towel": { # "cover with towel" task, 45 demos @ 5 Hz control "image_obs_keys": {"primary": "static_image", "secondary": None, "wrist": None}, "depth_obs_keys": {"primary": "static_depth_image", "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, @@ -645,29 +645,65 @@ class ActionEncoding(IntEnum): "libero_spatial_no_noops": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "libero_object_no_noops": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "libero_goal_no_noops": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, "libero_10_no_noops": { "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"}, "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, - "state_obs_keys": ["EEF_state", None, "gripper_state"], + "state_obs_keys": ["EEF_state", "gripper_state"], "state_encoding": StateEncoding.POS_EULER, "action_encoding": ActionEncoding.EEF_POS, }, + "libero_4_task_suites_no_noops": { + "image_obs_keys": {"primary": "image", "secondary": None, "wrist": "wrist_image"}, + "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, + "state_obs_keys": ["EEF_state", "gripper_state"], + "state_encoding": StateEncoding.POS_EULER, + "action_encoding": ActionEncoding.EEF_POS, + }, + ### ALOHA fine-tuning datasets + "aloha1_fold_shorts_20_demos": { + "image_obs_keys": {"primary": "image", "secondary": None, "left_wrist": "left_wrist_image", "right_wrist": "right_wrist_image"}, + "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, + "state_obs_keys": ["state"], + "state_encoding": StateEncoding.JOINT_BIMANUAL, + "action_encoding": ActionEncoding.JOINT_POS_BIMANUAL, + }, + "aloha1_fold_shirt_30_demos": { + "image_obs_keys": {"primary": "image", "secondary": None, "left_wrist": "left_wrist_image", "right_wrist": "right_wrist_image"}, + "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, + "state_obs_keys": ["state"], + "state_encoding": StateEncoding.JOINT_BIMANUAL, + "action_encoding": ActionEncoding.JOINT_POS_BIMANUAL, + }, + "aloha1_scoop_X_into_bowl_45_demos": { + "image_obs_keys": {"primary": "image", "secondary": None, "left_wrist": "left_wrist_image", "right_wrist": "right_wrist_image"}, + "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, + "state_obs_keys": ["state"], + "state_encoding": StateEncoding.JOINT_BIMANUAL, + "action_encoding": ActionEncoding.JOINT_POS_BIMANUAL, + }, + "aloha1_put_X_into_pot_300_demos": { + "image_obs_keys": {"primary": "image", "secondary": None, "left_wrist": "left_wrist_image", "right_wrist": "right_wrist_image"}, + "depth_obs_keys": {"primary": None, "secondary": None, "wrist": None}, + "state_obs_keys": ["state"], + "state_encoding": StateEncoding.JOINT_BIMANUAL, + "action_encoding": ActionEncoding.JOINT_POS_BIMANUAL, + }, } diff --git a/prismatic/vla/datasets/rlds/oxe/materialize.py b/prismatic/vla/datasets/rlds/oxe/materialize.py index 56d0d38fe..fd4103d8d 100644 --- a/prismatic/vla/datasets/rlds/oxe/materialize.py +++ b/prismatic/vla/datasets/rlds/oxe/materialize.py @@ -10,9 +10,9 @@ from typing import Any, Dict, List, Tuple from prismatic.overwatch import initialize_overwatch +from prismatic.vla.constants import ACTION_DIM, ACTION_PROPRIO_NORMALIZATION_TYPE, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX, NUM_ACTIONS_CHUNK, PROPRIO_DIM, STOP_INDEX from prismatic.vla.datasets.rlds.oxe.configs import OXE_DATASET_CONFIGS, ActionEncoding from prismatic.vla.datasets.rlds.oxe.transforms import OXE_STANDARDIZATION_TRANSFORMS -from prismatic.vla.datasets.rlds.utils.data_utils import NormalizationType # Initialize Overwatch =>> Wraps `logging.Logger` overwatch = initialize_overwatch(__name__) @@ -25,12 +25,12 @@ def make_oxe_dataset_kwargs( load_depth: bool = False, load_proprio: bool = True, load_language: bool = True, - action_proprio_normalization_type: NormalizationType = NormalizationType.NORMAL, + action_proprio_normalization_type = ACTION_PROPRIO_NORMALIZATION_TYPE, ) -> Dict[str, Any]: """Generates config (kwargs) for given dataset from Open-X Embodiment.""" dataset_kwargs = deepcopy(OXE_DATASET_CONFIGS[dataset_name]) - if dataset_kwargs["action_encoding"] not in [ActionEncoding.EEF_POS, ActionEncoding.EEF_R6]: - raise ValueError(f"Cannot load `{dataset_name}`; only EEF_POS & EEF_R6 actions supported!") + if dataset_kwargs["action_encoding"] not in [ActionEncoding.EEF_POS, ActionEncoding.EEF_R6, ActionEncoding.JOINT_POS_BIMANUAL]: + raise ValueError(f"Cannot load `{dataset_name}`; only EEF_POS & EEF_R6 & JOINT_POS_BIMANUAL actions supported!") # [Contract] For EEF_POS & EEF_R6 actions, only the last action dimension (gripper) is absolute! # Normalize all action dimensions *except* the gripper @@ -40,6 +40,9 @@ def make_oxe_dataset_kwargs( elif dataset_kwargs["action_encoding"] is ActionEncoding.EEF_R6: dataset_kwargs["absolute_action_mask"] = [False] * 9 + [True] dataset_kwargs["action_normalization_mask"] = [True] * 9 + [False] + elif dataset_kwargs["action_encoding"] is ActionEncoding.JOINT_POS_BIMANUAL: + dataset_kwargs["absolute_action_mask"] = [True] * 14 + dataset_kwargs["action_normalization_mask"] = [True] * 14 dataset_kwargs["action_proprio_normalization_type"] = action_proprio_normalization_type # Adjust Loaded Camera Views @@ -83,7 +86,7 @@ def get_oxe_dataset_kwargs_and_weights( load_depth: bool = False, load_proprio: bool = True, load_language: bool = True, - action_proprio_normalization_type: NormalizationType = NormalizationType.NORMAL, + action_proprio_normalization_type = ACTION_PROPRIO_NORMALIZATION_TYPE, ) -> Tuple[Dict[str, Any], List[float]]: """ Generates dataset kwargs for a given dataset mix from the Open X-Embodiment dataset. The returned kwargs diff --git a/prismatic/vla/datasets/rlds/oxe/mixtures.py b/prismatic/vla/datasets/rlds/oxe/mixtures.py index aca03da44..c5a2862fd 100644 --- a/prismatic/vla/datasets/rlds/oxe/mixtures.py +++ b/prismatic/vla/datasets/rlds/oxe/mixtures.py @@ -206,5 +206,25 @@ "libero_10_no_noops": [ ("libero_10_no_noops", 1.0), ], -} + "libero_4_task_suites_no_noops": [ + ("libero_spatial_no_noops", 1.0), + ("libero_object_no_noops", 1.0), + ("libero_goal_no_noops", 1.0), + ("libero_10_no_noops", 1.0), + ], + + # === ALOHA Fine-Tuning Datasets === + "aloha1_fold_shorts_20_demos": [ + ("aloha1_fold_shorts_20_demos", 1.0), + ], + "aloha1_fold_shirt_30_demos": [ + ("aloha1_fold_shirt_30_demos", 1.0), + ], + "aloha1_scoop_X_into_bowl_45_demos": [ + ("aloha1_scoop_X_into_bowl_45_demos", 1.0), + ], + "aloha1_put_X_into_pot_300_demos": [ + ("aloha1_put_X_into_pot_300_demos", 1.0), + ], # fmt: on +} diff --git a/prismatic/vla/datasets/rlds/oxe/transforms.py b/prismatic/vla/datasets/rlds/oxe/transforms.py index cc9c68712..bf848e98f 100644 --- a/prismatic/vla/datasets/rlds/oxe/transforms.py +++ b/prismatic/vla/datasets/rlds/oxe/transforms.py @@ -841,6 +841,11 @@ def libero_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]: return trajectory +def aloha_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]: + # Don't need to do anything because dataset is already in the correct format + return trajectory + + # === Registry === OXE_STANDARDIZATION_TRANSFORMS = { "bridge_oxe": bridge_oxe_dataset_transform, @@ -919,4 +924,10 @@ def libero_dataset_transform(trajectory: Dict[str, Any]) -> Dict[str, Any]: "libero_object_no_noops": libero_dataset_transform, "libero_goal_no_noops": libero_dataset_transform, "libero_10_no_noops": libero_dataset_transform, + "libero_4_task_suites_no_noops": libero_dataset_transform, + ### ALOHA fine-tuning datasets + "aloha1_fold_shorts_20_demos": aloha_dataset_transform, + "aloha1_fold_shirt_30_demos": aloha_dataset_transform, + "aloha1_scoop_X_into_bowl_45_demos": aloha_dataset_transform, + "aloha1_put_X_into_pot_300_demos": aloha_dataset_transform, } diff --git a/prismatic/vla/datasets/rlds/traj_transforms.py b/prismatic/vla/datasets/rlds/traj_transforms.py index 82cc43a2e..d2ae695ab 100644 --- a/prismatic/vla/datasets/rlds/traj_transforms.py +++ b/prismatic/vla/datasets/rlds/traj_transforms.py @@ -24,24 +24,22 @@ def chunk_act_obs(traj: Dict, window_size: int, future_action_window_size: int = """ traj_len = tf.shape(traj["action"])[0] action_dim = traj["action"].shape[-1] - chunk_indices = tf.broadcast_to(tf.range(-window_size + 1, 1), [traj_len, window_size]) + tf.broadcast_to( - tf.range(traj_len)[:, None], [traj_len, window_size] + effective_traj_len = traj_len - future_action_window_size + chunk_indices = tf.broadcast_to(tf.range(-window_size + 1, 1), [effective_traj_len, window_size]) + tf.broadcast_to( + tf.range(effective_traj_len)[:, None], [effective_traj_len, window_size] ) action_chunk_indices = tf.broadcast_to( tf.range(-window_size + 1, 1 + future_action_window_size), - [traj_len, window_size + future_action_window_size], + [effective_traj_len, window_size + future_action_window_size], ) + tf.broadcast_to( - tf.range(traj_len)[:, None], - [traj_len, window_size + future_action_window_size], + tf.range(effective_traj_len)[:, None], + [effective_traj_len, window_size + future_action_window_size], ) floored_chunk_indices = tf.maximum(chunk_indices, 0) - if "timestep" in traj["task"]: - goal_timestep = traj["task"]["timestep"] - else: - goal_timestep = tf.fill([traj_len], traj_len - 1) + goal_timestep = tf.fill([effective_traj_len], traj_len - 1) floored_action_chunk_indices = tf.minimum(tf.maximum(action_chunk_indices, 0), goal_timestep[:, None]) @@ -51,22 +49,10 @@ def chunk_act_obs(traj: Dict, window_size: int, future_action_window_size: int = # indicates whether an entire observation is padding traj["observation"]["pad_mask"] = chunk_indices >= 0 - # if no absolute_action_mask was provided, assume all actions are relative - if "absolute_action_mask" not in traj and future_action_window_size > 0: - logging.warning( - "future_action_window_size > 0 but no absolute_action_mask was provided. " - "Assuming all actions are relative for the purpose of making neutral actions." - ) - absolute_action_mask = traj.get("absolute_action_mask", tf.zeros([traj_len, action_dim], dtype=tf.bool)) - neutral_actions = tf.where( - absolute_action_mask[:, None, :], - traj["action"], # absolute actions are repeated (already done during chunking) - tf.zeros_like(traj["action"]), # relative actions are zeroed - ) - - # actions past the goal timestep become neutral - action_past_goal = action_chunk_indices > goal_timestep[:, None] - traj["action"] = tf.where(action_past_goal[:, :, None], neutral_actions, traj["action"]) + # Truncate other elements of the trajectory dict + traj["task"] = tf.nest.map_structure(lambda x: tf.gather(x, tf.range(effective_traj_len)), traj["task"]) + traj["dataset_name"] = tf.gather(traj["dataset_name"], tf.range(effective_traj_len)) + traj["absolute_action_mask"] = tf.gather(traj["absolute_action_mask"], tf.range(effective_traj_len)) return traj diff --git a/prismatic/vla/datasets/rlds/utils/data_utils.py b/prismatic/vla/datasets/rlds/utils/data_utils.py index 7b0e5ae9c..41b61bd12 100644 --- a/prismatic/vla/datasets/rlds/utils/data_utils.py +++ b/prismatic/vla/datasets/rlds/utils/data_utils.py @@ -7,7 +7,6 @@ import hashlib import json import os -from enum import Enum from typing import Any, Callable, Dict, List, Optional, Tuple import dlimp as dl @@ -16,6 +15,7 @@ from tqdm import tqdm from prismatic.overwatch import initialize_overwatch +from prismatic.vla.constants import NormalizationType # Initialize Overwatch =>> Wraps `logging.Logger` overwatch = initialize_overwatch(__name__) @@ -45,15 +45,6 @@ def to_padding(tensor: tf.Tensor) -> tf.Tensor: raise ValueError(f"Cannot generate padding for tensor of type {tensor.dtype}.") -# Defines supported normalization schemes for action and proprioceptive state. -class NormalizationType(str, Enum): - # fmt: off - NORMAL = "normal" # Normalize to Mean = 0, Stdev = 1 - BOUNDS = "bounds" # Normalize to Interval = [-1, 1] - BOUNDS_Q99 = "bounds_q99" # Normalize [quantile_01, ..., quantile_99] --> [-1, ..., 1] - # fmt: on - - # === State / Action Processing Primitives === diff --git a/pyproject.toml b/pyproject.toml index f72cae071..562e9ba27 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,17 +3,17 @@ requires = ["setuptools"] build-backend = "setuptools.build_meta" [project] -name = "openvla" +name = "openvla-oft" authors = [ {name = "Moo Jin Kim", email="moojink@stanford.edu"}, - {name = "Karl Pertsch", email="pertsch@berkeley.edu"}, - {name = "Siddharth Karamcheti", email="skaramcheti@cs.stanford.edu"}, + {name = "Chelsea Finn", email="cbfinn@cs.stanford.edu"}, + {name = "Percy Liang", email="pliang@cs.stanford.edu"}, ] -description = "OpenVLA: Vision-Language-Action Models for Robotics" -version = "0.0.3" +description = "Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success" +version = "0.0.1" readme = "README.md" requires-python = ">=3.8" -keywords = ["vision-language-actions models", "multimodal pretraining", "robot learning"] +keywords = ["vision-language-actions models", "fine-tuning", "robot learning"] license = {file = "LICENSE"} classifiers = [ "Development Status :: 3 - Alpha", @@ -47,12 +47,17 @@ dependencies = [ "torch==2.2.0", "torchvision==0.17.0", "torchaudio==2.2.0", - "transformers==4.40.1", + "transformers @ git+https://github.com/moojink/transformers-openvla-oft.git", # IMPORTANT: Use this fork for bidirectional attn (for parallel decoding) "wandb", "tensorflow==2.15.0", "tensorflow_datasets==4.9.3", "tensorflow_graphics==2021.12.3", - "dlimp @ git+https://github.com/moojink/dlimp_openvla" + "dlimp @ git+https://github.com/moojink/dlimp_openvla", + "diffusers", + "imageio", + "uvicorn", + "fastapi", + "json-numpy", ] [project.optional-dependencies] @@ -69,9 +74,9 @@ sagemaker = [ ] [project.urls] -homepage = "https://github.com/openvla/openvla" -repository = "https://github.com/openvla/openvla" -documentation = "https://github.com/openvla/openvla" +homepage = "https://github.com/moojink/openvla-oft" +repository = "https://github.com/moojink/openvla-oft" +documentation = "https://github.com/moojink/openvla-oft" [tool.setuptools.packages.find] where = ["."] diff --git a/third_party/lerobot b/third_party/lerobot new file mode 160000 index 000000000..aca464ca7 --- /dev/null +++ b/third_party/lerobot @@ -0,0 +1 @@ +Subproject commit aca464ca72aba644b918d403f7e76e39bfed3317 diff --git a/vla-scripts/deploy.py b/vla-scripts/deploy.py index c70a9f279..def1bc373 100644 --- a/vla-scripts/deploy.py +++ b/vla-scripts/deploy.py @@ -1,30 +1,7 @@ """ deploy.py -Provide a lightweight server/client implementation for deploying OpenVLA models (through the HF AutoClass API) over a -REST API. This script implements *just* the server, with specific dependencies and instructions below. - -Note that for the *client*, usage just requires numpy/json-numpy, and requests; example usage below! - -Dependencies: - => Server (runs OpenVLA model on GPU): `pip install uvicorn fastapi json-numpy` - => Client: `pip install requests json-numpy` - -Client (Standalone) Usage (assuming a server running on 0.0.0.0:8000): - -``` -import requests -import json_numpy -json_numpy.patch() -import numpy as np - -action = requests.post( - "http://0.0.0.0:8000/act", - json={"image": np.zeros((256, 256, 3), dtype=np.uint8), "instruction": "do something"} -).json() - -Note that if your server is not accessible on the open web, you can use ngrok, or forward ports to your client via ssh: - => `ssh -L 8000:localhost:8000 ssh USER@` +Starts VLA server which the client can query to get robot actions. """ import os.path @@ -35,6 +12,7 @@ json_numpy.patch() import json import logging +import numpy as np import traceback from dataclasses import dataclass from pathlib import Path @@ -48,61 +26,69 @@ from PIL import Image from transformers import AutoModelForVision2Seq, AutoProcessor -# === Utilities === -SYSTEM_PROMPT = ( - "A chat between a curious user and an artificial intelligence assistant. " - "The assistant gives helpful, detailed, and polite answers to the user's questions." +from experiments.robot.openvla_utils import ( + get_vla, + get_vla_action, + get_action_head, + get_processor, + get_proprio_projector, +) +from experiments.robot.robot_utils import ( + get_image_resize_size, ) +from prismatic.vla.constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX, NUM_ACTIONS_CHUNK, PROPRIO_DIM, STOP_INDEX def get_openvla_prompt(instruction: str, openvla_path: Union[str, Path]) -> str: - if "v01" in openvla_path: - return f"{SYSTEM_PROMPT} USER: What action should the robot take to {instruction.lower()}? ASSISTANT:" - else: - return f"In: What action should the robot take to {instruction.lower()}?\nOut:" + return f"In: What action should the robot take to {instruction.lower()}?\nOut:" # === Server Interface === class OpenVLAServer: - def __init__(self, openvla_path: Union[str, Path], attn_implementation: Optional[str] = "flash_attention_2") -> Path: + def __init__(self, cfg) -> Path: """ - A simple server for OpenVLA models; exposes `/act` to predict an action for a given image + instruction. - => Takes in {"image": np.ndarray, "instruction": str, "unnorm_key": Optional[str]} - => Returns {"action": np.ndarray} + A simple server for OpenVLA models; exposes `/act` to predict an action for a given observation + instruction. """ - self.openvla_path, self.attn_implementation = openvla_path, attn_implementation - self.device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu") - - # Load VLA Model using HF AutoClasses - self.processor = AutoProcessor.from_pretrained(self.openvla_path, trust_remote_code=True) - self.vla = AutoModelForVision2Seq.from_pretrained( - self.openvla_path, - attn_implementation=attn_implementation, - torch_dtype=torch.bfloat16, - low_cpu_mem_usage=True, - trust_remote_code=True, - ).to(self.device) - - # [Hacky] Load Dataset Statistics from Disk (if passing a path to a fine-tuned model) - if os.path.isdir(self.openvla_path): - with open(Path(self.openvla_path) / "dataset_statistics.json", "r") as f: - self.vla.norm_stats = json.load(f) - - def predict_action(self, payload: Dict[str, Any]) -> str: + self.cfg = cfg + + # Load model + self.vla = get_vla(cfg) + + # Load proprio projector + self.proprio_projector = None + if cfg.use_proprio: + self.proprio_projector = get_proprio_projector(cfg, self.vla.llm_dim, PROPRIO_DIM) + + # Load continuous action head + self.action_head = None + if cfg.use_l1_regression or cfg.use_diffusion: + self.action_head = get_action_head(cfg, self.vla.llm_dim) + + # Check that the model contains the action un-normalization key + assert cfg.unnorm_key in self.vla.norm_stats, f"Action un-norm key {cfg.unnorm_key} not found in VLA `norm_stats`!" + + # Get Hugging Face processor + self.processor = None + self.processor = get_processor(cfg) + + # Get expected image dimensions + self.resize_size = get_image_resize_size(cfg) + + + def get_server_action(self, payload: Dict[str, Any]) -> str: try: if double_encode := "encoded" in payload: # Support cases where `json_numpy` is hard to install, and numpy arrays are "double-encoded" as strings assert len(payload.keys()) == 1, "Only uses encoded payload!" payload = json.loads(payload["encoded"]) - # Parse payload components - image, instruction = payload["image"], payload["instruction"] - unnorm_key = payload.get("unnorm_key", None) + observation = payload + instruction = observation["instruction"] + + action = get_vla_action( + self.cfg, self.vla, self.processor, observation, instruction, action_head=self.action_head, proprio_projector=self.proprio_projector, use_film=self.cfg.use_film, + ) - # Run VLA Inference - prompt = get_openvla_prompt(instruction, self.openvla_path) - inputs = self.processor(prompt, Image.fromarray(image).convert("RGB")).to(self.device, dtype=torch.bfloat16) - action = self.vla.predict_action(**inputs, unnorm_key=unnorm_key, do_sample=False) if double_encode: return JSONResponse(json_numpy.dumps(action)) else: @@ -111,33 +97,56 @@ def predict_action(self, payload: Dict[str, Any]) -> str: logging.error(traceback.format_exc()) logging.warning( "Your request threw an error; make sure your request complies with the expected format:\n" - "{'image': np.ndarray, 'instruction': str}\n" - "You can optionally an `unnorm_key: str` to specific the dataset statistics you want to use for " - "de-normalizing the output actions." + "{'observation': dict, 'instruction': str}\n" ) return "error" - def run(self, host: str = "0.0.0.0", port: int = 8000) -> None: + def run(self, host: str = "0.0.0.0", port: int = 8777) -> None: self.app = FastAPI() - self.app.post("/act")(self.predict_action) + self.app.post("/act")(self.get_server_action) uvicorn.run(self.app, host=host, port=port) @dataclass class DeployConfig: # fmt: off - openvla_path: Union[str, Path] = "openvla/openvla-7b" # HF Hub Path (or path to local run directory) # Server Configuration host: str = "0.0.0.0" # Host IP Address - port: int = 8000 # Host Port - + port: int = 8777 # Host Port + + ################################################################################################################# + # Model-specific parameters + ################################################################################################################# + model_family: str = "openvla" # Model family + pretrained_checkpoint: Union[str, Path] = "" # Pretrained checkpoint path + + use_l1_regression: bool = True # If True, uses continuous action head with L1 regression objective + use_diffusion: bool = False # If True, uses continuous action head with diffusion modeling objective (DDIM) + num_diffusion_steps: int = 50 # (When `diffusion==True`) Number of diffusion steps for inference + use_film: bool = False # If True, uses FiLM to infuse language inputs into visual features + num_images_in_input: int = 3 # Number of images in the VLA input (default: 3) + use_proprio: bool = True # Whether to include proprio state in input + + center_crop: bool = True # Center crop? (if trained w/ random crop image aug) + num_open_loop_steps: int = 25 # Number of actions to execute open-loop before requerying policy + + unnorm_key: Union[str, Path] = "" # Action un-normalization key + use_relative_actions: bool = False # Whether to use relative actions (delta joint angles) + + load_in_8bit: bool = False # (For OpenVLA only) Load with 8-bit quantization + load_in_4bit: bool = False # (For OpenVLA only) Load with 4-bit quantization + + ################################################################################################################# + # Utils + ################################################################################################################# + seed: int = 7 # Random Seed (for reproducibility) # fmt: on @draccus.wrap() def deploy(cfg: DeployConfig) -> None: - server = OpenVLAServer(cfg.openvla_path) + server = OpenVLAServer(cfg) server.run(cfg.host, port=cfg.port) diff --git a/vla-scripts/finetune.py b/vla-scripts/finetune.py index ec51a6b3c..f5c1146c6 100644 --- a/vla-scripts/finetune.py +++ b/vla-scripts/finetune.py @@ -1,175 +1,878 @@ """ finetune.py -Simple script for parameter-efficient fine-tuning of OpenVLA models loaded through the HuggingFace AutoClasses, using -HuggingFace PEFT library for low-rank adaptation (LoRA). - -Notes & Benchmarks: - - Requires PEFT (`pip install peft==0.11.1`) - - LoRA fine-tuning (see parameters below -- no quantization, LoRA rank = 32, target_modules = all-linear): - + One 48 GB GPU can fit a Batch Size of 12 - + One 80 GB GPU can fit a Batch Size of 24 - -Run with: - - [Single Node Multi-GPU (= $K) ]: torchrun --standalone --nnodes 1 --nproc-per-node $K vla-scripts/finetune.py - - [Override Config Values]: torchrun --standalone --nnodes 1 --nproc-per-node $K vla-scripts/finetune.py \ - --data_root_dir \ - --dataset_name \ - --run_root_dir \ - ... +Fine-tunes OpenVLA via LoRA. """ import os +import time +import json from collections import deque from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import Dict, Optional, Tuple, Type import draccus +from lerobot.common.datasets.factory import make_dataset import torch import torch.distributed as dist +import torch.nn as nn import tqdm from accelerate import PartialState -from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training +from huggingface_hub import HfApi, snapshot_download +from peft import LoraConfig, PeftModel, get_peft_model 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 torch.optim.lr_scheduler import MultiStepLR +from torch.utils.data import DataLoader, RandomSampler +from transformers import ( + AutoConfig, + AutoImageProcessor, + AutoModelForVision2Seq, + AutoProcessor, +) 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 experiments.robot.openvla_utils import ( + check_model_logic_mismatch, + model_is_on_hf_hub, + update_auto_map, +) + +from lerobot.common.datasets.lerobot_dataset import LeRobotDataset 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 +from prismatic.models.action_heads import DiffusionActionHead, L1RegressionActionHead +from prismatic.models.backbones.llm.prompting import PurePromptBuilder +from prismatic.models.film_vit_wrapper import FiLMedPrismaticVisionBackbone +from prismatic.models.projectors import ( + NoisyActionProjector, + ProprioProjector, +) +from prismatic.training.train_utils import ( + compute_actions_l1_loss, + compute_token_accuracy, + get_current_action_mask, + get_next_actions_mask, +) +from prismatic.util.extern.hf.lerobot_utils import ( + create_action_norm_stats_dict_from_lerobot_dataset, + create_rlds_dataset_stats_dict_from_lerobot_dataset, + VLACollatorForLeRobotDataset, +) +from prismatic.vla.action_tokenizer import ActionTokenizer +from prismatic.vla.constants import ( + ACTION_DIM, + ACTION_PROPRIO_NORMALIZATION_TYPE, + NUM_ACTIONS_CHUNK, + PROPRIO_DIM, +) +from prismatic.vla.datasets import ( + RLDSBatchTransform, + RLDSDataset, +) +from prismatic.vla.datasets.rlds.utils.data_utils import save_dataset_statistics +from prismatic.util.data_utils import greyscale_float_tensor_preprocessing_wrapper + # Sane Defaults os.environ["TOKENIZERS_PARALLELISM"] = "false" -# # === Utilities === -# # fmt: off -# def create_vision_transform(vla: nn.Module, input_size: int) -> Callable[[Image.Image], torch.Tensor]: -# """Gets image transform for the vision encoder.""" -# data_cfg = timm.data.resolve_model_data_config(vla.vision_backbone) -# data_cfg["input_size"] = (3, input_size, input_size) -# return timm.data.create_transform( -# input_size=data_cfg["input_size"], -# interpolation=data_cfg["interpolation"], -# mean=data_cfg["mean"], -# std=data_cfg["std"], -# crop_pct=1.0, # Set to 1.0 to disable cropping -# crop_mode="center", # Default crop mode --> no-op when `crop_pct == 1.0` -# is_training=False, # Disable image_aug when loading transform; handled by RLDS dataloader -# ) -# -# # fmt: on - - @dataclass class FinetuneConfig: # fmt: off - vla_path: str = "openvla/openvla-7b" # Path to OpenVLA model (on HuggingFace Hub) - - # Directory Paths - data_root_dir: Path = Path("datasets/open-x-embodiment") # Path to Open-X dataset directory - dataset_name: str = "droid_wipe" # Name of fine-tuning dataset (e.g., `droid_wipe`) - run_root_dir: Path = Path("runs") # Path to directory to store logs & checkpoints - adapter_tmp_dir: Path = Path("adapter-tmp") # Temporary directory for LoRA weights before fusing - - # Fine-tuning Parameters - batch_size: int = 16 # Fine-tuning batch size - max_steps: int = 200_000 # Max number of fine-tuning steps - save_steps: int = 5000 # Interval for checkpoint saving - learning_rate: float = 5e-4 # Fine-tuning learning rate - grad_accumulation_steps: int = 1 # Gradient accumulation steps - image_aug: bool = True # Whether to train with image augmentations - shuffle_buffer_size: int = 100_000 # Dataloader shuffle buffer size (can reduce if OOM) - save_latest_checkpoint_only: bool = True # Whether to save only one checkpoint per run and - # continually overwrite the latest checkpoint - # (If False, saves all checkpoints) - - # LoRA Arguments - use_lora: bool = True # Whether to use LoRA fine-tuning - lora_rank: int = 32 # Rank of LoRA weight matrix - lora_dropout: float = 0.0 # Dropout applied to LoRA weights - use_quantization: bool = False # Whether to 4-bit quantize VLA for LoRA fine-tuning - # => CAUTION: Reduces memory but hurts performance - - # Tracking Parameters - wandb_project: str = "openvla" # Name of W&B project to log to (use default!) - wandb_entity: str = "stanford-voltron" # Name of entity to log under - run_id_note: Optional[str] = None # Extra note for logging, Weights & Biases + vla_path: str = "openvla/openvla-7b" # Path to OpenVLA model (on HuggingFace Hub or stored locally) + + # Dataset + data_root_dir: Path = Path("datasets/rlds") # Directory containing RLDS datasets + dataset_name: str = "aloha_scoop_x_into_bowl" # Name of fine-tuning dataset (e.g., `aloha_scoop_x_into_bowl`) + run_root_dir: Path = Path("runs") # Path to directory to store logs & checkpoints + shuffle_buffer_size: int = 100_000 # Dataloader shuffle buffer size (can reduce if OOM errors occur) + + # Algorithm and architecture + use_l1_regression: bool = True # If True, trains continuous action head with L1 regression objective + use_diffusion: bool = False # If True, trains continuous action head with diffusion modeling objective (DDIM) + num_diffusion_steps: int = 50 # (When `diffusion==True`) Number of diffusion steps for training + use_film: bool = False # If True, uses FiLM to infuse language inputs into visual features + num_images_in_input: int = 1 # Number of images in the VLA input (default: 1) + use_proprio: bool = False # If True, includes robot proprioceptive state in input + + # Training configuration + batch_size: int = 8 # Batch size per device (total batch size = batch_size * num GPUs) + learning_rate: float = 5e-4 # Learning rate + lr_warmup_steps: int = 0 # Number of steps to warm up learning rate (from 10% to 100%) + num_steps_before_decay: int = 100_000 # Number of steps before LR decays by 10x + grad_accumulation_steps: int = 1 # Number of gradient accumulation steps + max_steps: int = 200_000 # Max number of training steps + use_val_set: bool = False # If True, uses validation set and log validation metrics + val_freq: int = 10_000 # (When `use_val_set==True`) Validation set logging frequency in steps + val_time_limit: int = 180 # (When `use_val_set==True`) Time limit for computing validation metrics + save_freq: int = 10_000 # Checkpoint saving frequency in steps + save_latest_checkpoint_only: bool = False # If True, saves only 1 checkpoint, overwriting latest checkpoint + # (If False, saves all checkpoints) + resume: bool = False # If True, resumes from checkpoint + resume_step: Optional[int] = None # (When `resume==True`) Step number that we are resuming from + image_aug: bool = True # If True, trains with image augmentations (HIGHLY RECOMMENDED) + diffusion_sample_freq: int = 50 # (When `use_diffusion==True`) Frequency for sampling in steps + + # LoRA + use_lora: bool = True # If True, uses LoRA fine-tuning + lora_rank: int = 32 # Rank of LoRA weight matrix + lora_dropout: float = 0.0 # Dropout applied to LoRA weights + merge_lora_during_training: bool = True # If True, merges LoRA weights and saves result during training + # Note: Merging can be very slow on some machines. If so, set to + # False and merge final checkpoint offline! + + # Logging + wandb_entity: str = "your-wandb-entity" # Name of WandB entity + wandb_project: str = "your-wandb-project" # Name of WandB project + run_id_note: Optional[str] = None # Extra note to add to end of run ID for logging + run_id_override: Optional[str] = None # Optional string to override the run ID with + wandb_log_freq: int = 10 # WandB logging frequency in steps + + # LeRobot dataset + use_lerobot_dataset: bool = False + lerobot_dataset_root_dir: Path = Path("data") + lerobot_dataset_name: str = "robotgeneralist/nomagic-simple-box" + lerobot_tolerance_s: float = 0.01 # fmt: on + + + +def remove_ddp_in_checkpoint(state_dict) -> dict: + """ + Removes the 'module.' prefix from parameter names in a PyTorch model state dictionary that was saved using + DistributedDataParallel (DDP). + + When a model is trained using PyTorch's DistributedDataParallel, the saved state dictionary contains parameters + prefixed with 'module.'. This function removes these prefixes to make the state dictionary compatible when + loading into models that are not yet wrapped in DDP. + + Args: + state_dict (dict): PyTorch model state dictionary. + + Returns: + dict: A new state dictionary with the same contents but with 'module.' prefixes removed from parameter names. + Parameters without the 'module.' prefix remain unchanged. + """ + new_state_dict = {} + for k, v in state_dict.items(): + if k[:7] == "module.": + new_state_dict[k[7:]] = v + else: + new_state_dict[k] = v + return new_state_dict + + +def get_run_id(cfg) -> str: + """ + Generates or retrieves an identifier string for an experiment run. + + Args: + cfg (FinetuneConfig): Training configuration. + + Returns: + str: Experiment run ID. + """ + if cfg.run_id_override is not None: + # Override the run ID with the user-provided ID + run_id = cfg.run_id_override + elif cfg.resume: + # Override run ID with the previous resumed run's ID + run_id = cfg.vla_path.split("/")[-1] + # Remove the "--XXX_chkpt" suffix from the run ID if it exists + if "chkpt" in run_id.split("--")[-1]: + run_id = "--".join(run_id.split("--")[:-1]) + else: + run_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: + run_id += f"+lora-r{cfg.lora_rank}+dropout-{cfg.lora_dropout}" + if cfg.image_aug: + run_id += "--image_aug" + if cfg.run_id_note is not None: + run_id += f"--{cfg.run_id_note}" + return run_id + + +def load_checkpoint(module_name: str, path: str, step: int, device: str = "cpu") -> dict: + """ + Loads a checkpoint for a given module. + + Args: + module_name (str): Name of model component to load checkpoint for. + path (str): Path to checkpoint directory. + step (int): Gradient step number of saved checkpoint. + device (str): String specifying how to remap storage locations (default = "cpu"). + + Returns: + dict: PyTorch model state dictionary. + """ + checkpoint_path = os.path.join(path, f"{module_name}--{step}_checkpoint.pt") + print(f"Loading checkpoint: {checkpoint_path}") + state_dict = torch.load(checkpoint_path, weights_only=True, map_location=device) + return remove_ddp_in_checkpoint(state_dict) + + +def wrap_ddp(module: nn.Module, device_id: int, find_unused: bool = False) -> DDP: + """ + Wrap a module with DistributedDataParallel. + + Args: + module (nn.Module): PyTorch module. + device_id (str): Device ID. + find_unused (bool): Whether to detect parameters without gradients in distributed training. + + Returns: + DistributedDataParallel: PyTorch module wrapped with DDP. + """ + return DDP(module, device_ids=[device_id], find_unused_parameters=find_unused, gradient_as_bucket_view=True) + + +def count_parameters(module: nn.Module, name: str) -> None: + """ + Counts and prints the number of trainable parameters in a module. + + Args: + module (nn.Module): PyTorch module. + module_name (str): Name of model component. + + Returns: + None. + """ + num_params = sum(p.numel() for p in module.parameters() if p.requires_grad) + print(f"# trainable params in {name}: {num_params}") + + +def init_module( + module_class: Type[nn.Module], + module_name: str, + cfg: FinetuneConfig, + device_id: int, + module_args: dict, + to_bf16: bool = False, + find_unused_params: bool = False, +) -> DDP: + """ + Initializes a module, optionally loads checkpoint, moves to device, and wraps with DDP. + + Args: + module_class (Type[nn.Module]): Class of PyTorch module to initialize. + module_name (str): Name of model component to load checkpoint for. + cfg (FinetuneConfig): Training configuration. + device_id (str): Device ID. + module_args (dict): Args for initializing the module. + to_bf16 (bool): Whether to convert to torch.bfloat16 data type. + find_unused_params (bool): Whether to detect parameters without gradients in distributed training. + + Returns: + DistributedDataParallel: PyTorch module wrapped with DDP. + """ + module = module_class(**module_args) + count_parameters(module, module_name) + + if cfg.resume: + state_dict = load_checkpoint(module_name, cfg.vla_path, cfg.resume_step) + module.load_state_dict(state_dict) + + if to_bf16: + module = module.to(torch.bfloat16) + module = module.to(device_id) + + return wrap_ddp(module, device_id, find_unused_params) + + +def run_forward_pass( + vla, + action_head, + noisy_action_projector, + proprio_projector, + batch, + action_tokenizer, + device_id, + use_l1_regression, + use_diffusion, + use_proprio, + use_film, + num_patches, + compute_diffusion_l1=False, + num_diffusion_steps=None, +) -> Tuple[torch.Tensor, Dict[str, float]]: + """ + Compute model forward pass and metrics for both training and validation. + + Args: + vla (OpenVLAForActionPrediction): Vision-language-action policy. + action_head (nn.Module): Action head module. + noisy_action_projector (nn.Module): Noisy action projector module (only used for diffusion). + proprio_projector (nn.Module): Proprioceptive state projector module. + batch (dict): Input batch. + action_tokenizer (ActionTokenizer): Action tokenizer. + device_id (str): Device ID. + use_l1_regression (bool): Whether to use L1 regression. + use_diffusion (bool): Whether to use diffusion. + use_proprio (bool): Whether to use proprioceptive state as input. + use_film (bool): Whether to use FiLM for better language following. + num_patches (int): Number of vision patches. + compute_diffusion_l1 (bool): Whether to sample actions and compute L1 loss for diffusion (do this once every + diffusion_sample_freq steps during training; do it every batch for validation) + num_diffusion_steps (int): Number of diffusion steps (only used for diffusion). + + Returns: + tuple: (loss, metrics_dict) + loss: The loss tensor with gradient for backpropagation. + metrics_dict: Dictionary of computed metrics (detached values for logging). + """ + metrics = {} + + # Get ground-truth action labels + ground_truth_actions = batch["actions"].to(device_id).to(torch.bfloat16) + + # [Only for diffusion] Sample noisy actions used as input for noise predictor network + if use_diffusion: + noisy_dict = action_head.module.sample_noisy_actions(ground_truth_actions) + noise, noisy_actions, diffusion_timestep_embeddings = ( + noisy_dict["noise"], + noisy_dict["noisy_actions"], + noisy_dict["diffusion_timestep_embeddings"], + ) + else: + noise, noisy_actions, diffusion_timestep_embeddings = None, None, None + + # VLA forward pass + with torch.autocast("cuda", dtype=torch.bfloat16): + output: CausalLMOutputWithPast = vla( + input_ids=batch["input_ids"].to(device_id), + attention_mask=batch["attention_mask"].to(device_id), + pixel_values=batch["pixel_values"].to(torch.bfloat16).to(device_id), + labels=batch["labels"], + output_hidden_states=True, + proprio=batch["proprio"] if use_proprio else None, + proprio_projector=proprio_projector if use_proprio else None, + noisy_actions=noisy_actions if use_diffusion else None, + noisy_action_projector=noisy_action_projector if use_diffusion else None, + diffusion_timestep_embeddings=diffusion_timestep_embeddings if use_diffusion else None, + use_film=use_film, + ) + + # Get action masks needed for logging + ground_truth_token_ids = batch["labels"][:, 1:].to(device_id) + current_action_mask = get_current_action_mask(ground_truth_token_ids) + next_actions_mask = get_next_actions_mask(ground_truth_token_ids) + + # Compute metrics for discrete action representation (next-token prediction) + if not (use_l1_regression or use_diffusion): + loss = output.loss + predicted_token_ids = output.logits[:, num_patches:-1].argmax(dim=2) + curr_action_accuracy = compute_token_accuracy( + predicted_token_ids, ground_truth_token_ids, mask=current_action_mask + ) + curr_action_l1_loss = compute_actions_l1_loss( + action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask=current_action_mask + ) + next_actions_accuracy = compute_token_accuracy( + predicted_token_ids, ground_truth_token_ids, mask=next_actions_mask + ) + next_actions_l1_loss = compute_actions_l1_loss( + action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask=next_actions_mask + ) + metrics.update( + { + "loss_value": loss.item(), # Detached value for logging + "curr_action_accuracy": curr_action_accuracy.item(), + "curr_action_l1_loss": curr_action_l1_loss.item(), + "next_actions_accuracy": next_actions_accuracy.item(), + "next_actions_l1_loss": next_actions_l1_loss.item(), + } + ) + # Compute metrics for continuous action representations (L1 regression | diffusion) + else: + # Get last layer hidden states + last_hidden_states = output.hidden_states[-1] # (B, seq_len, D) + # Get hidden states for text portion of prompt+response (after the vision patches) + text_hidden_states = last_hidden_states[:, num_patches:-1] + # Get hidden states for action portion of response + batch_size = batch["input_ids"].shape[0] + actions_hidden_states = ( + text_hidden_states[current_action_mask | next_actions_mask] + .reshape(batch_size, NUM_ACTIONS_CHUNK * ACTION_DIM, -1) + .to(torch.bfloat16) + ) # (B, act_chunk_len, D) + + if use_l1_regression: + # Predict action + predicted_actions = action_head.module.predict_action(actions_hidden_states) + # Get full L1 loss + loss = torch.nn.L1Loss()(ground_truth_actions, predicted_actions) + + if use_diffusion: + # Predict noise + noise_pred = action_head.module.predict_noise(actions_hidden_states) + # Get diffusion noise prediction MSE loss + noise_pred = noise_pred.reshape(noise.shape) + loss = nn.functional.mse_loss(noise_pred, noise, reduction="mean") + + # Only sample actions and compute L1 losses if specified + if compute_diffusion_l1: + with torch.no_grad(): + predicted_actions = run_diffusion_sampling( + vla=vla, + action_head=action_head, + noisy_action_projector=noisy_action_projector, + proprio_projector=proprio_projector, + batch=batch, + batch_size=batch_size, + num_patches=num_patches, + actions_shape=ground_truth_actions.shape, + device_id=device_id, + current_action_mask=current_action_mask, + next_actions_mask=next_actions_mask, + use_proprio=use_proprio, + use_film=use_film, + ) + + metrics.update( + { + "loss_value": loss.item(), # Detached value for logging + } + ) + + # Get detailed L1 losses for logging + should_log_l1_loss = not use_diffusion or (use_diffusion and compute_diffusion_l1) + if should_log_l1_loss: + ground_truth_curr_action = ground_truth_actions[:, 0] + predicted_curr_action = predicted_actions[:, 0] + ground_truth_next_actions = ground_truth_actions[:, 1:] + predicted_next_actions = predicted_actions[:, 1:] + curr_action_l1_loss = torch.nn.L1Loss()(ground_truth_curr_action, predicted_curr_action) + next_actions_l1_loss = torch.nn.L1Loss()(ground_truth_next_actions, predicted_next_actions) + metrics.update( + { + "curr_action_l1_loss": curr_action_l1_loss.item(), + "next_actions_l1_loss": next_actions_l1_loss.item(), + } + ) + + # Return both the loss tensor (with gradients) and the metrics dictionary (with detached values) + return loss, metrics + + +def run_diffusion_sampling( + vla, + action_head, + noisy_action_projector, + proprio_projector, + batch, + batch_size, + num_patches, + actions_shape, + device_id, + current_action_mask, + next_actions_mask, + use_proprio, + use_film, +) -> torch.Tensor: + """ + Run diffusion sampling (reverse diffusion) to generate actions. + + Args: + vla (OpenVLAForActionPrediction): Vision-language-action policy. + action_head (nn.Module): Action head module. + noisy_action_projector (nn.Module): Noisy action projector module (only used for diffusion). + proprio_projector (nn.Module): Proprioceptive state projector module. + batch (dict): Input batch. + batch_size (int): Batch size. + num_patches (int): Number of vision patches. + actions_shape (tuple): Shape of ground-truth actions. + device_id (str): Device ID. + current_action_mask (torch.Tensor): Mask for current action. + next_actions_mask (torch.Tensor): Mask for next actions. + use_proprio (bool): Whether to use proprioceptive state as input. + use_film (bool): Whether to use FiLM for better language following. + + Returns: + torch.Tensor: Predicted actions. + """ + # Sample random noisy action, used as the starting point for reverse diffusion + noise = torch.randn( + size=(batch_size, NUM_ACTIONS_CHUNK, ACTION_DIM), + device=device_id, + dtype=torch.bfloat16, + ) # (B, chunk_len, action_dim) + + # Set diffusion timestep values + action_head.module.noise_scheduler.set_timesteps(action_head.module.num_diffusion_steps) + + # Reverse diffusion: Iteratively denoise to generate action, conditioned on observation + curr_noisy_actions = noise + for t in action_head.module.noise_scheduler.timesteps: + # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action embedding, + # and diffusion timestep embedding) + timesteps = torch.Tensor([t]).repeat(batch_size).to(device_id) + diffusion_timestep_embeddings = ( + action_head.module.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device) + ) # (B, llm_dim) + diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1) # (B, 1, llm_dim) + + with torch.autocast("cuda", dtype=torch.bfloat16): + output = vla( + input_ids=batch["input_ids"].to(device_id), + attention_mask=batch["attention_mask"].to(device_id), + pixel_values=batch["pixel_values"].to(torch.bfloat16).to(device_id), + labels=batch["labels"], + output_hidden_states=True, + proprio=batch["proprio"] if use_proprio else None, + proprio_projector=proprio_projector if use_proprio else None, + noisy_actions=curr_noisy_actions, + noisy_action_projector=noisy_action_projector, + diffusion_timestep_embeddings=diffusion_timestep_embeddings, + use_film=use_film, + ) + # Get last layer hidden states + last_hidden_states = output.hidden_states[-1] # (B, seq_len, D) + # Get hidden states for text portion of prompt+response (after the vision patches) + text_hidden_states = last_hidden_states[:, num_patches:-1] + # Get hidden states for action portion of response + actions_hidden_states = text_hidden_states[current_action_mask | next_actions_mask].reshape( + batch_size, NUM_ACTIONS_CHUNK * ACTION_DIM, -1 + ) # (B, act_chunk_len, D) + actions_hidden_states = actions_hidden_states.to(torch.bfloat16) + # Predict noise + noise_pred = action_head.module.predict_noise(actions_hidden_states) + + # Compute the action at the previous diffusion timestep: x_t -> x_{t-1} + curr_noisy_actions = action_head.module.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample + + return curr_noisy_actions.reshape(actions_shape) + + +def compute_smoothened_metrics(metrics_deques) -> dict: + """ + Compute smoothened metrics from recent deques. + + Args: + metrics_deques (dict): Dictionary of deques containing recent metrics. + + Returns: + dict: Dictionary of smoothened metrics. + """ + smoothened_metrics = {} + for name, deque in metrics_deques.items(): + if deque and len(deque) > 0: + smoothened_metrics[name] = sum(deque) / len(deque) + return smoothened_metrics + + +def log_metrics_to_wandb(metrics, prefix, step, wandb_entity) -> None: + """ + Log metrics to Weights & Biases. + + Args: + metrics (dict): Dictionary of metrics to log + prefix (str): Prefix for metric names + step (int): Training step + wandb_entity (str): W&B entity instance + + Returns: + None. + """ + log_dict = {} + for name, value in metrics.items(): + # Map loss_value to Loss for better readability in W&B + if name == "loss_value": + log_dict[f"{prefix}/Loss"] = value + # Keep other metrics as is + else: + log_dict[f"{prefix}/{name.replace('_', ' ').title()}"] = value + wandb_entity.log(log_dict, step=step) + + +def save_training_checkpoint( + cfg, + run_dir, + log_step, + vla, + processor, + proprio_projector, + noisy_action_projector, + action_head, + train_dataset, + distributed_state, +) -> None: + """ + Save all training checkpoints including model components, LoRA adapter, and dataset statistics. + + Args: + cfg (FinetuneConfig): Training configuration. + run_dir (Path): Experiment run directory path. + log_step (int): Current logging step. + vla (OpenVLAForActionPrediction): Vision-language-action policy. + processor (PrismaticProcessor): OpenVLA inputs processor. + proprio_projector (nn.Module): Proprioceptive state projector module. + noisy_action_projector (nn.Module): Noisy action projector module (only used for diffusion). + action_head (nn.Module): Action head module. + train_dataset (RLDSDataset): Training dataset. + distributed_state (PartialState): Distributed training state. + + Returns: + None. + """ + # Determine checkpoint paths and naming + if cfg.save_latest_checkpoint_only: + checkpoint_dir = run_dir + checkpoint_name_suffix = "latest_checkpoint.pt" + else: + checkpoint_dir = Path(str(run_dir) + f"--{log_step}_chkpt") + checkpoint_name_suffix = f"{log_step}_checkpoint.pt" + + adapter_dir = checkpoint_dir / "lora_adapter" + + # Create directories and save dataset statistics (main process only) + if distributed_state.is_main_process: + os.makedirs(checkpoint_dir, exist_ok=True) + os.makedirs(adapter_dir, exist_ok=True) + dataset_stats \ + = create_rlds_dataset_stats_dict_from_lerobot_dataset( + train_dataset, + dataset_name="train" + ) + save_dataset_statistics(dataset_stats, checkpoint_dir) + print(f"Saving Model Checkpoint for Step {log_step}") + + # Wait for directories to be created + dist.barrier() + + # Save model components (main process only) + if distributed_state.is_main_process: + # Save processor and LoRA adapter + processor.save_pretrained(checkpoint_dir) + vla.module.save_pretrained(adapter_dir) + + # Save other components + if cfg.use_proprio and proprio_projector is not None: + torch.save(proprio_projector.state_dict(), checkpoint_dir / f"proprio_projector--{checkpoint_name_suffix}") + + if cfg.use_diffusion and noisy_action_projector is not None: + torch.save( + noisy_action_projector.state_dict(), checkpoint_dir / f"noisy_action_projector--{checkpoint_name_suffix}" + ) + + if (cfg.use_l1_regression or cfg.use_diffusion) and action_head is not None: + torch.save(action_head.state_dict(), checkpoint_dir / f"action_head--{checkpoint_name_suffix}") + + if cfg.use_film: + # To be safe, just save the entire vision backbone (not just FiLM components) + torch.save( + vla.module.vision_backbone.state_dict(), checkpoint_dir / f"vision_backbone--{checkpoint_name_suffix}" + ) + + # Wait for model components to be saved + dist.barrier() + + # Merge LoRA weights into base model and save resulting model checkpoint + # Note: Can be very slow on some devices; if so, we recommend merging offline + if cfg.use_lora and cfg.merge_lora_during_training: + 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 distributed_state.is_main_process: + merged_vla.save_pretrained(checkpoint_dir) + print(f"Saved merged model for Step {log_step} at: {checkpoint_dir}") + + # Wait for merged model to be saved + dist.barrier() + + +def run_validation( + vla, + action_head, + noisy_action_projector, + proprio_projector, + val_dataloader, + action_tokenizer, + device_id, + cfg, + num_patches, + log_step, + distributed_state, + val_time_limit, +) -> None: + """ + Compute validation set metrics for logging. + + Args: + vla (OpenVLAForActionPrediction): Vision-language-action policy. + action_head (nn.Module): Action head module. + noisy_action_projector (nn.Module): Noisy action projector module (only used for diffusion). + proprio_projector (nn.Module): Proprioceptive state projector module. + val_dataloader (DataLoader): Validation data loader. + action_tokenizer (ActionTokenizer): Action tokenizer. + device_id (str): Device ID. + cfg (FinetuneConfig): Training configuration. + num_patches (int): Number of vision patches. + log_step (int): Current logging step. + distributed_state (PartialState): Distributed training state. + val_time_limit (int): Time limit for computing validation metrics. + + Returns: + None. + """ + val_start_time = time.time() + vla.eval() + val_batches_count = 0 + + # List to store validation metrics + all_val_metrics = [] + + with torch.no_grad(): + for batch in val_dataloader: + # Always compute L1 loss for validation, even for diffusion + _, metrics = run_forward_pass( + vla=vla, + action_head=action_head, + noisy_action_projector=noisy_action_projector, + proprio_projector=proprio_projector, + batch=batch, + action_tokenizer=action_tokenizer, + device_id=device_id, + use_l1_regression=cfg.use_l1_regression, + use_diffusion=cfg.use_diffusion, + use_proprio=cfg.use_proprio, + use_film=cfg.use_film, + num_patches=num_patches, + compute_diffusion_l1=True, + num_diffusion_steps=cfg.num_diffusion_steps if cfg.use_diffusion else None, + ) + + # Add the loss value to the metrics + metrics["loss"] = metrics["loss_value"] + all_val_metrics.append(metrics) + val_batches_count += 1 + + # Cut testing on validation set short if it exceeds time limit + if time.time() - val_start_time > val_time_limit: + break + + # Compute average validation metrics + avg_val_metrics = {} + for metric_name in all_val_metrics[0].keys(): + values = [metrics[metric_name] for metrics in all_val_metrics if metric_name in metrics] + if values: + avg_val_metrics[metric_name] = sum(values) / len(values) + + # Add batch count to metrics + avg_val_metrics["val_batches_count"] = val_batches_count + + # Log validation metrics to W&B + if distributed_state.is_main_process: + log_metrics_to_wandb(avg_val_metrics, "VLA Val", log_step, wandb) @draccus.wrap() def finetune(cfg: FinetuneConfig) -> None: + """ + Fine-tunes base VLA on demonstration dataset via LoRA. + + Allows toggling different action representations (discrete vs. continuous), different learning objectives + (next-token prediction vs. L1 regression vs. diffusion), FiLM. Also allows for additional model inputs, + such as additional camera images and robot proprioceptive state. Assumes parallel action generation with + action chunking. + + Args: + cfg (FinetuneConfig): Training configuration. + + Returns: + None. + """ + assert cfg.use_lora, "Only LoRA fine-tuning is supported. Please set --use_lora=True!" + assert not (cfg.use_l1_regression and cfg.use_diffusion), ( + "Cannot do both L1 regression and diffusion. Please pick one of them!" + ) + + # Trim trailing forward slash ('/') in VLA path if it exists + cfg.vla_path = cfg.vla_path.rstrip("/") print(f"Fine-tuning OpenVLA Model `{cfg.vla_path}` on `{cfg.dataset_name}`") - # [Validate] Ensure GPU Available & Set Device / Distributed Context - assert torch.cuda.is_available(), "Fine-tuning assumes at least one GPU is available!" + # Get experiment run ID + run_id = get_run_id(cfg) + + # Create experiment run directory + run_dir = cfg.run_root_dir / run_id + os.makedirs(run_dir, exist_ok=True) + + # GPU setup distributed_state = PartialState() - torch.cuda.set_device(device_id := distributed_state.local_process_index) + device_id = distributed_state.local_process_index + torch.cuda.set_device(device_id) torch.cuda.empty_cache() - # Configure Unique Experiment ID & Log Directory - 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}" + # Initialize wandb logging + if distributed_state.is_main_process: + wandb.init(entity=cfg.wandb_entity, project=cfg.wandb_project, name=f"ft+{run_id}") + + # Print detected constants + print( + "Detected constants:\n" + f"\tNUM_ACTIONS_CHUNK: {NUM_ACTIONS_CHUNK}\n" + f"\tACTION_DIM: {ACTION_DIM}\n" + f"\tPROPRIO_DIM: {PROPRIO_DIM}\n" + f"\tACTION_PROPRIO_NORMALIZATION_TYPE: {ACTION_PROPRIO_NORMALIZATION_TYPE}" ) - 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 is not None: - exp_id += f"--{cfg.run_id_note}" - if cfg.image_aug: - exp_id += "--image_aug" - - # Start =>> Build Directories - 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 =>> only if LoRA fine-tuning - quantization_config = None - if cfg.use_quantization: - assert cfg.use_lora, "Quantized training only supported for LoRA fine-tuning!" - quantization_config = BitsAndBytesConfig( - load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_type="nf4" - ) + # Two options: + # (1) Base model is on Hugging Face Hub + # - Then download it and record the path to the download directory + # (2) Base model is stored locally + # - Then register model config in HF Auto Classes + # In both cases, we want to check whether any changes have been made to + # the `modeling_prismatic.py` file in this codebase; if so, we will copy + # the file to the downloaded or locally stored checkpoint directory so + # that the user's changes to the VLA class logic go into effect + if model_is_on_hf_hub(cfg.vla_path): + # Download model directly from Hugging Face Hub + vla_download_path = snapshot_download(repo_id=cfg.vla_path) + # Overwrite VLA path + cfg.vla_path = vla_download_path + else: + # Register OpenVLA model to HF Auto Classes (not needed if the model is on HF Hub) + AutoConfig.register("openvla", OpenVLAConfig) + AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) + AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) + AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) + + # Update config.json and sync model files + if distributed_state.is_main_process: + update_auto_map(cfg.vla_path) + check_model_logic_mismatch(cfg.vla_path) - # Register OpenVLA model to HF Auto Classes (not needed if the model is on HF Hub) - AutoConfig.register("openvla", OpenVLAConfig) - AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) - AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) - AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) + # Wait for model files to be synced + dist.barrier() - # Load OpenVLA Processor and Model using HF AutoClasses + # Load processor and VLA 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, low_cpu_mem_usage=True, trust_remote_code=True, - ) + ).to(device_id) - # Device Placement =>> note that BitsAndBytes automatically handles for quantized training - if cfg.use_quantization: - vla = prepare_model_for_kbit_training(vla) - else: - vla = vla.to(device_id) + # Set number of images in VLA input + vla.vision_backbone.set_num_images_in_input(cfg.num_images_in_input) - # [LoRA] Wrap Model w/ PEFT `LoraConfig` =>> by default we set `target_modules=all-linear` + # LoRA setup if cfg.use_lora: lora_config = LoraConfig( r=cfg.lora_rank, @@ -181,193 +884,364 @@ def finetune(cfg: FinetuneConfig) -> None: vla = get_peft_model(vla, lora_config) vla.print_trainable_parameters() - # Wrap VLA in PyTorch DDP Wrapper for Multi-GPU Training - vla = DDP(vla, device_ids=[device_id], find_unused_parameters=True, gradient_as_bucket_view=True) + # FiLM setup + if cfg.use_film: + count_parameters(vla.vision_backbone, "vla.vision_backbone (original)") + # Wrap vision backbone with FiLM wrapper + # Important: For this, must specify `vla.model.vision_backbone` instead of just `vla.vision_backbone`, since the + # latter would cause the new wrapped backbone to be saved as a new attribute of `vla` instead of overwriting the + # original one (due to the LoRA wrapper) + vla.model.vision_backbone = FiLMedPrismaticVisionBackbone( + vision_backbone=vla.model.vision_backbone, + llm_dim=vla.llm_dim, + ) + count_parameters(vla.vision_backbone, "vla.vision_backbone (post-wrap)") + if cfg.resume: + state_dict = load_checkpoint("vision_backbone", cfg.vla_path, cfg.resume_step) + vla.model.vision_backbone.load_state_dict(state_dict) + vla.model.vision_backbone = vla.model.vision_backbone.to(device_id) + + # Wrap VLA with DDP + vla = wrap_ddp(vla, device_id, find_unused=True) + + # If applicable, instantiate proprio projector + if cfg.use_proprio: + proprio_projector = init_module( + ProprioProjector, + "proprio_projector", + cfg, + device_id, + {"llm_dim": vla.module.llm_dim, "proprio_dim": PROPRIO_DIM}, + ) + + # If applicable, instantiate continuous action head for L1 regression + if cfg.use_l1_regression: + action_head = init_module( + L1RegressionActionHead, + "action_head", + cfg, + device_id, + {"input_dim": vla.module.llm_dim, "hidden_dim": vla.module.llm_dim, "action_dim": ACTION_DIM}, + to_bf16=True, + ) - # Create Optimizer =>> note that we default to a simple constant learning rate! + # If applicable, instantiate diffusion action head and noisy action projector + if cfg.use_diffusion: + action_head = init_module( + DiffusionActionHead, + "action_head", + cfg, + device_id, + { + "input_dim": vla.module.llm_dim, + "hidden_dim": vla.module.llm_dim, + "action_dim": ACTION_DIM, + "num_diffusion_steps": cfg.num_diffusion_steps, + }, + to_bf16=True, + ) + noisy_action_projector = init_module( + NoisyActionProjector, "noisy_action_projector", cfg, device_id, {"llm_dim": vla.module.llm_dim} + ) + + # Get number of vision patches + NUM_PATCHES = vla.module.vision_backbone.get_num_patches() * vla.module.vision_backbone.get_num_images_in_input() + # If we have proprio inputs, a single proprio embedding is appended to the end of the vision patch embeddings + if cfg.use_proprio: + NUM_PATCHES += 1 + # For diffusion, a single diffusion timestep embedding is appended to the end of the vision patch embeddings + if cfg.use_diffusion: + NUM_PATCHES += 1 + + # Instantiate optimizer trainable_params = [param for param in vla.parameters() if param.requires_grad] + if cfg.use_l1_regression or cfg.use_diffusion: + trainable_params += [param for param in action_head.parameters() if param.requires_grad] + if cfg.use_diffusion: + trainable_params += [param for param in noisy_action_projector.parameters() if param.requires_grad] + if cfg.use_proprio: + trainable_params += [param for param in proprio_projector.parameters() if param.requires_grad] + print(f"# total trainable params: {sum(p.numel() for p in trainable_params)}") optimizer = AdamW(trainable_params, lr=cfg.learning_rate) + # Record original learning rate + original_lr = optimizer.param_groups[0]["lr"] + + # Create learning rate scheduler + scheduler = MultiStepLR( + optimizer, + milestones=[cfg.num_steps_before_decay], # Number of steps after which LR will change + gamma=0.1, # Multiplicative factor of learning rate decay + ) + # Create Action Tokenizer action_tokenizer = ActionTokenizer(processor.tokenizer) + # TODO: FIgure out what ActionTokenizer and processor.tokenizer do + # TODO: and how to duplicate this functionality in LeRobotDataset # Load Fine-tuning Dataset =>> note that we use an RLDS-formatted dataset following Open X-Embodiment by default. # =>> If you want to use a non-RLDS dataset (e.g., a standard PyTorch Dataset) see the following commented block. # =>> Note that our training code does not loop over epochs because the RLDS loader does this implicitly; if using # your own Dataset, make sure to add the appropriate logic to the training loop! - # + # # TODO: Figure this out # --- - # from prismatic.vla.datasets import DummyDataset - # - # vla_dataset = DummyDataset( + + if not cfg.use_lerobot_dataset: + raise NotImplementedError("Only LeRobotDataset is currently supported") + + import numpy as np + if cfg.use_lerobot_dataset: + repo_dir = Path(cfg.lerobot_dataset_root_dir) / cfg.lerobot_dataset_name + + # Wrap the image transform function to handle tensors + wrapped_transform = greyscale_float_tensor_preprocessing_wrapper( + processor.image_processor.apply_transform + ) + + train_dataset = LeRobotDataset( + repo_id="NotRequired", + root=repo_dir, + episodes=None, + image_transforms=wrapped_transform, + delta_timestamps=None, + tolerance_s=cfg.lerobot_tolerance_s, + download_videos=False, + local_files_only=True, + video_backend=None, + ) + + + # 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, + # prompt_builder_fn=PurePromptBuilder, + # ) + # vla_dataset = RLDSDataset( + # cfg.data_root_dir, + # cfg.dataset_name, + # batch_transform, + # resize_resolution=tuple(vla.module.config.image_sizes), + # shuffle_buffer_size=cfg.shuffle_buffer_size, + # image_aug=cfg.image_aug, # ) - # --- - 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.module.config.image_sizes), - shuffle_buffer_size=cfg.shuffle_buffer_size, - image_aug=cfg.image_aug, - ) - - # [Important] Save Dataset Statistics =>> used to de-normalize actions for inference! - if distributed_state.is_main_process: - save_dataset_statistics(vla_dataset.dataset_statistics, run_dir) - # Create Collator and DataLoader - collator = PaddedCollatorForActionPrediction( - processor.tokenizer.model_max_length, processor.tokenizer.pad_token_id, padding_side="right" - ) - dataloader = DataLoader( - vla_dataset, - batch_size=cfg.batch_size, - sampler=None, - collate_fn=collator, - num_workers=0, # Important =>> Set to 0 if using RLDS; TFDS rolls its own parallelism! - ) + # Create collator and dataloader + if cfg.use_lerobot_dataset: + dataset_name = cfg.lerobot_dataset_name.split("/")[-1] + + action_norm_stats \ + = create_action_norm_stats_dict_from_lerobot_dataset(train_dataset) + dataset_stats \ + = create_rlds_dataset_stats_dict_from_lerobot_dataset( + train_dataset, + dataset_name="train", + ) + + # Save dataset statistics for unnorming actions during inference + if distributed_state.is_main_process: + save_dataset_statistics(dataset_stats, run_dir) + + vla_collator = VLACollatorForLeRobotDataset( + action_tokenizer=action_tokenizer, + base_tokenizer=processor.tokenizer, + prompt_builder_fn=PurePromptBuilder, + pad_token_id=processor.tokenizer.pad_token_id, + model_max_length=processor.tokenizer.model_max_length, + predict_stop_token=True, + use_wrist_image=cfg.num_images_in_input > 1, + use_proprio=cfg.use_proprio, + action_norm_stats=action_norm_stats, + ) + + if cfg.use_val_set: + train_subset, val_subset \ + = create_train_val_split_from_lerobot_dataset( + train_dataset, + split=0.1, + ) + + train_indices_list = [int(idx) for idx in train_subset.indices] + val_indices_list = [int(idx) for idx in val_subset.indices] - # Initialize Logging =>> W&B - if distributed_state.is_main_process: - wandb.init(entity=cfg.wandb_entity, project=cfg.wandb_project, name=f"ft+{exp_id}") + print(f"All train indices: {train_indices_list}") + print(f"All val indices: {val_indices_list}") + + train_sampler = RandomSampler(train_subset) + dataloader = DataLoader( + train_subset, + batch_size=cfg.batch_size, + sampler=train_sampler, + collate_fn=vla_collator, + num_workers=4, + pin_memory=True, + ) + val_dataloader = DataLoader( + val_subset, + batch_size=cfg.batch_size, + shuffle=False, + collate_fn=vla_collator, + num_workers=4, + pin_memory=True, + ) + else: + sampler = RandomSampler(train_dataset) + dataloader = DataLoader( + train_dataset, + batch_size=cfg.batch_size, + sampler=sampler, + collate_fn=vla_collator, + num_workers=4, + pin_memory=True, + ) # Deque to store recent train metrics (used for computing smoothened metrics for gradient accumulation) - 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) + recent_metrics = { + "loss_value": deque(maxlen=cfg.grad_accumulation_steps), + "curr_action_accuracy": deque(maxlen=cfg.grad_accumulation_steps), + "curr_action_l1_loss": deque(maxlen=cfg.grad_accumulation_steps), + "next_actions_accuracy": deque(maxlen=cfg.grad_accumulation_steps), + "next_actions_l1_loss": deque(maxlen=cfg.grad_accumulation_steps), + } - # Train! + # Start training 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_id), - attention_mask=batch["attention_mask"].to(device_id), - pixel_values=batch["pixel_values"].to(torch.bfloat16).to(device_id), - labels=batch["labels"], - ) - loss = output.loss - - # Normalize loss to account for gradient accumulation - normalized_loss = loss / cfg.grad_accumulation_steps + + # Compute the number of epochs needed to train for cfg.max_steps steps + # =>> This is used to set the number of epochs in the progress bar + steps_per_epoch = len(dataloader) # number of batches per epoch + min_epochs = (cfg.max_steps * cfg.grad_accumulation_steps) // steps_per_epoch + 1 + num_epochs = min_epochs - # Backward pass - normalized_loss.backward() - - # Compute Accuracy and L1 Loss for Logging - action_logits = output.logits[:, vla.module.vision_backbone.featurizer.patch_embed.num_patches : -1] - action_preds = action_logits.argmax(dim=2) - action_gt = batch["labels"][:, 1:].to(action_preds.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()) - ) - continuous_actions_gt = torch.tensor( - action_tokenizer.decode_token_ids_to_actions(action_gt[mask].cpu().numpy()) - ) - action_l1_loss = torch.nn.functional.l1_loss(continuous_actions_pred, continuous_actions_gt) - - # Store recent train 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 - - # Compute smoothened train metrics - # =>> Equal to current step metrics when not using gradient accumulation - # =>> Otherwise, equal to the average of metrics observed over micro-batches used for gradient accumulation - smoothened_loss = sum(recent_losses) / len(recent_losses) - smoothened_action_accuracy = sum(recent_action_accuracies) / len(recent_action_accuracies) - smoothened_l1_loss = sum(recent_l1_losses) / len(recent_l1_losses) - - # Push Metrics to W&B (every 10 gradient steps) - if distributed_state.is_main_process and gradient_step_idx % 10 == 0: - wandb.log( - { - "train_loss": smoothened_loss, - "action_accuracy": smoothened_action_accuracy, - "l1_loss": smoothened_l1_loss, - }, - step=gradient_step_idx, + total_gradient_step_idx = 0 + for _ in range(num_epochs): + for batch_idx, batch in enumerate(dataloader): + # Compute training metrics and loss + compute_diffusion_l1 = cfg.use_diffusion and batch_idx % cfg.diffusion_sample_freq == 0 + loss, metrics = run_forward_pass( + vla=vla, + action_head=action_head, + noisy_action_projector=noisy_action_projector if cfg.use_diffusion else None, + proprio_projector=proprio_projector if cfg.use_proprio else None, + batch=batch, + action_tokenizer=action_tokenizer, + device_id=device_id, + use_l1_regression=cfg.use_l1_regression, + use_diffusion=cfg.use_diffusion, + use_proprio=cfg.use_proprio, + use_film=cfg.use_film, + num_patches=NUM_PATCHES, + compute_diffusion_l1=compute_diffusion_l1, + num_diffusion_steps=cfg.num_diffusion_steps if cfg.use_diffusion else None, ) - # Optimizer Step - if (batch_idx + 1) % cfg.grad_accumulation_steps == 0: - optimizer.step() - optimizer.zero_grad() - progress.update() - - # Save Model Checkpoint =>> by default, only keeps the latest checkpoint, continually overwriting it! - if gradient_step_idx > 0 and gradient_step_idx % cfg.save_steps == 0: - if distributed_state.is_main_process: - print(f"Saving Model Checkpoint for Step {gradient_step_idx}") + # Normalize loss to account for gradient accumulation + normalized_loss = loss / cfg.grad_accumulation_steps - # If LoRA, we first save adapter weights, then merge into full model; otherwise, default save! - save_dir = adapter_dir if cfg.use_lora else run_dir + # Backward pass + normalized_loss.backward() - # Save Processor & Weights - processor.save_pretrained(run_dir) - vla.module.save_pretrained(save_dir) + # Store recent train metrics + for metric_name, value in metrics.items(): + if metric_name in recent_metrics: + recent_metrics[metric_name].append(value) - # Wait for processor and adapter weights to be saved by main process - dist.barrier() + # Compute (per-epoch) gradient step index + epoch_gradient_step_idx = batch_idx // cfg.grad_accumulation_steps + + # Compute smoothened train metrics + smoothened_metrics = compute_smoothened_metrics(recent_metrics) - # Merge LoRA weights into model backbone for faster inference - # =>> Note that merging is slow and can be done post-hoc to speed up training - 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 distributed_state.is_main_process: - if cfg.save_latest_checkpoint_only: - # Overwrite latest checkpoint - merged_vla.save_pretrained(run_dir) - - print(f"Saved Model Checkpoint for Step {gradient_step_idx} at: {run_dir}") - else: - # Prepare to save checkpoint in new directory - checkpoint_dir = Path(str(run_dir) + f"--{gradient_step_idx}_chkpt") - os.makedirs(checkpoint_dir, exist_ok=True) + # Push Metrics to W&B (every wandb_log_freq gradient steps) + log_step = total_gradient_step_idx if not cfg.resume else cfg.resume_step + total_gradient_step_idx + if distributed_state.is_main_process and log_step % cfg.wandb_log_freq == 0: + log_metrics_to_wandb(smoothened_metrics, "VLA Train", log_step, wandb) - # Save dataset statistics to new directory - save_dataset_statistics(vla_dataset.dataset_statistics, checkpoint_dir) + # [If applicable] Linearly warm up learning rate from 10% to 100% of original + if cfg.lr_warmup_steps > 0: + lr_progress = min((epoch_gradient_step_idx + 1) / cfg.lr_warmup_steps, 1.0) # Cap at 1.0 + current_lr = original_lr * (0.1 + 0.9 * lr_progress) + for param_group in optimizer.param_groups: + param_group["lr"] = current_lr - # Save processor and model weights to new directory - processor.save_pretrained(checkpoint_dir) - merged_vla.save_pretrained(checkpoint_dir) + if distributed_state.is_main_process and epoch_gradient_step_idx % cfg.wandb_log_freq == 0: + # Log the learning rate + # Make sure to do this AFTER any learning rate modifications (e.g., warmup/decay) + wandb.log( + { + "VLA Train/Learning Rate": scheduler.get_last_lr()[0], + }, + step=log_step, + ) - print(f"Saved Model Checkpoint for Step {gradient_step_idx} at: {checkpoint_dir}") + # Optimizer Step + max_num_grad_acc_steps_done: bool \ + = (batch_idx + 1) % cfg.grad_accumulation_steps == 0 + all_batches_done: bool \ + = batch_idx == len(dataloader) - 1 + time_to_step: bool = max_num_grad_acc_steps_done or all_batches_done + if time_to_step: + optimizer.step() + scheduler.step() + optimizer.zero_grad() + progress.update() + total_gradient_step_idx += 1 - # Block on Main Process Checkpointing - dist.barrier() + # Save model checkpoint:o either keep latest checkpoint only or all checkpoints + if ( + time_to_step + and log_step > 0 + and log_step % cfg.save_freq == 0 + ): + save_training_checkpoint( + cfg=cfg, + run_dir=run_dir, + log_step=log_step, + vla=vla, + processor=processor, + proprio_projector=proprio_projector if cfg.use_proprio else None, + noisy_action_projector=noisy_action_projector if cfg.use_diffusion else None, + action_head=action_head if (cfg.use_l1_regression or cfg.use_diffusion) else None, + train_dataset=train_dataset, + distributed_state=distributed_state, + ) - # Stop training when max_steps is reached - if gradient_step_idx == cfg.max_steps: + # Test model on validation set + if ( + time_to_step + and cfg.use_val_set + and log_step > 0 + and log_step % cfg.val_freq == 0 + ): + run_validation( + vla=vla, + action_head=action_head, + noisy_action_projector=noisy_action_projector if cfg.use_diffusion else None, + proprio_projector=proprio_projector if cfg.use_proprio else None, + val_dataloader=val_dataloader, + action_tokenizer=action_tokenizer, + device_id=device_id, + cfg=cfg, + num_patches=NUM_PATCHES, + log_step=log_step, + distributed_state=distributed_state, + val_time_limit=cfg.val_time_limit, + ) + # Set model back to training mode after validation + vla.train() + + # Stop training when max_steps is reached + if total_gradient_step_idx == cfg.max_steps: + print(f"Max step {cfg.max_steps} reached! Stopping training...") + break + + if total_gradient_step_idx == cfg.max_steps: print(f"Max step {cfg.max_steps} reached! Stopping training...") break + if __name__ == "__main__": finetune() diff --git a/vla-scripts/finetune.sub b/vla-scripts/finetune.sub new file mode 100644 index 000000000..0d5178765 --- /dev/null +++ b/vla-scripts/finetune.sub @@ -0,0 +1,90 @@ +#!/bin/bash +#SBATCH -N 1 +#SBATCH -n 1 +#SBATCH -c 16 +#SBATCH --gres=gpu:2 +#SBATCH -t 48:00:00 +#SBATCH -p a100 +#SBATCH --mem=100G +#SBATCH -o .slurmlog/slurm-%j.out +#SBATCH -e .slurmlog/slurm-%j.err + +# --- Environment Setup --- + +module purge # Start with a clean environment. +module load cuda/11.8 # REPLACE with the CORRECT CUDA version! + +# Activate the virtual environment (created by create_env.sh). +source .venv/bin/activate + +# Install flash-attn *after* loading CUDA. +pip install "flash-attn==2.5.5" --no-build-isolation + +# Set environment variables for Weights & Biases logging +if [[ -z "${WANDB_API_KEY}" ]]; then + echo "Warning: WANDB_API_KEY not set in environment, reading from .env file" + if [ -f .env ]; then + export WANDB_API_KEY=$(grep WANDB_API_KEY .env | cut -d '=' -f2) + else + echo "Error: .env file not found" + exit 1 + fi +fi + +if [[ -z "${WANDB_MODE}" ]]; then + export WANDB_MODE="online" +fi + +# --- Run the Finetuning Script --- + +# Set PYTHONPATH and run the finetuning script using torchrun. +# Note: We use the absolute path to finetune.py to avoid issues with +# relative paths within the Slurm job. +PYTHONPATH="$PYTHONPATH:$(pwd)/lerobot" \ +mkdir -p "$(pwd)/.runs" && \ +mkdir -p "$(pwd)/.slurmlog" && \ +torchrun \ + --standalone \ + --nnodes 1 \ + --nproc-per-node 2 \ + "$(pwd)"/vla-scripts/finetune.py \ + --vla_path "openvla/openvla-7b" \ + --data_root_dir "data" \ + --dataset_name "nomagic-simple-box" \ + --run_root_dir ".runs/" \ + --shuffle_buffer_size 100000 \ + --use_l1_regression true \ + --use_diffusion false \ + --num_diffusion_steps 50 \ + --use_film true \ + --num_images_in_input 3 \ + --use_proprio false \ + --batch_size 4 \ + --learning_rate 5e-4 \ + --lr_warmup_steps 0 \ + --num_steps_before_decay 100000 \ + --grad_accumulation_steps 2 \ + --max_steps 200000 \ + --use_val_set true \ + --val_freq 1000 \ + --val_time_limit 180 \ + --save_freq 1000 \ + --save_latest_checkpoint_only true \ + --resume false \ + --image_aug true \ + --use_lora true \ + --lora_rank 32 \ + --lora_dropout 0.1 \ + --merge_lora_during_training false \ + --wandb_entity robotgeneralist \ + --wandb_project openvla \ + --wandb_log_freq 100 \ + --use_lerobot_dataset true \ + --lerobot_dataset_root_dir "data" \ + --lerobot_dataset_name "robotgeneralist/nomagic-simple-box-action-chunk-size-8" \ + --lerobot_tolerance_s 0.01 \ + --constants_config ur5e + +# To run this script, edit the options above, and then +# execute the following command from the root repository directory: +# sbatch vla-scripts/finetune.sub \ No newline at end of file diff --git a/vla-scripts/merge_lora_weights_and_save.py b/vla-scripts/merge_lora_weights_and_save.py new file mode 100644 index 000000000..8c38c10e9 --- /dev/null +++ b/vla-scripts/merge_lora_weights_and_save.py @@ -0,0 +1,73 @@ +""" +Loads a checkpoint that only has a LoRA adapter (no merged model) and merges the adapter +into the base OpenVLA model. Saves the final checkpoint in the same directory. + +Make sure to specify the correct base checkpoint when running this script. For example, +- if you fine-tuned the default OpenVLA-7B model without modifications, then `--base_checkpoint=="openvla/openvla-7b"` +- if you fine-tuned a different model or resumed fine-tuning from a different checkpoint, then specify that base checkpoint +- if you fine-tuned the default OpenVLA-7B model with modifications to `modeling_prismatic.py` (OpenVLA class definition), + then the base checkpoint path should point to the checkpoint containing the modifications + +Usage: + python vla-scripts/merge_lora_weights_and_save.py \ + --base_checkpoint openvla/openvla-7b \ + --lora_finetuned_checkpoint_dir /PATH/TO/CHECKPOINT/DIR/ +""" + +import os +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Union + +import draccus +import torch +from peft import PeftModel +from transformers import AutoConfig, AutoImageProcessor, AutoModelForVision2Seq, AutoProcessor + +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 + + +@dataclass +class ConvertConfig: + # fmt: off + + base_checkpoint: Union[str, Path] = "" # Base model checkpoint path/dir (either openvla/openvla-7b or whichever model you fine-tuned / resumed training from) + lora_finetuned_checkpoint_dir: Union[str, Path] = "" # Checkpoint directory containing the LoRA adapter + + # fmt: on + + +@draccus.wrap() +def main(cfg: ConvertConfig) -> None: + # Register OpenVLA model to HF Auto Classes (not needed if the model is on HF Hub) + AutoConfig.register("openvla", OpenVLAConfig) + AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) + AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) + AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) + + # Load Model using HF AutoClasses + print(f"Loading base model: {cfg.base_checkpoint}") + vla = AutoModelForVision2Seq.from_pretrained( + cfg.base_checkpoint, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + trust_remote_code=True, + ) + + # Load LoRA weights and merge into base model, then save final checkpoint + print("Merging LoRA weights into base model...") + start_time = time.time() + merged_vla = PeftModel.from_pretrained(vla, os.path.join(cfg.lora_finetuned_checkpoint_dir, "lora_adapter")).to( + "cuda" + ) + merged_vla = merged_vla.merge_and_unload() + merged_vla.save_pretrained(cfg.lora_finetuned_checkpoint_dir) + print(f"\nMerging complete! Time elapsed (sec): {time.time() - start_time}") + print(f"\nSaved merged model checkpoint at:\n{cfg.lora_finetuned_checkpoint_dir}") + + +if __name__ == "__main__": + main()