[Paper Review] OpenVLA: AnOpen-Source Vision-Language-Action Model

https://arxiv.org/pdf/2406.09246

[ Abstract ]

Internet-scale의 VLM(비전-언어) 데이터와 다양한 로봇 데이터를 결합하여 Pretrained된 대규모 Policies은 로봇에게 새로운 기술을 가르치는 방식을 변화시킬 잠재력을 가지고 있다. 새로운 Task를 처음부터 훈련하는 대신, 이러한 VLA(비전-언어-행동) 모델을 Finetuning하여 성능이 좋은 일반화 가능한 Policies을 얻을 수 있다. 그러나 VLA에는 몇 가지 한계점이 존재한다.

1) 기존 VLA는 대부분 폐쇄적이고 공개적으로 접근하기 어렵고
2) 새로운 작업에 대해 VLA를 효율적으로 Finetuning하는 방법이 충분히 연구되지 않았다.

이러한 과제를 해결하기 위해 본 논문은 OpenVLA를 공개한다. OpenVLA는 970,000개의 실제 로봇 시연 데이터를 다양하게 수집한 데이터셋에서 학습된 70억 매개변수의 오픈소스 VLA이다. OpenVLA는 Llama 2 언어 모델과 DINOv2 및 SigLIP에서 Pretrain된 특징을 융합하는 Vision 인코더를 기반으로 구축되었다. 추가된 데이터 다양성과 새로운 모델 아키텍처 덕분에 OpenVLA는 RT-2-X (55B)와 같은 폐쇄형 모델을 능가하며, 29개의 작업과 여러 로봇에서 Absolute Task 성공률이 16.5% 더 높고, Parameter가 7배 더 적다.

또한 Diffusion Policy와 같은 Scratch Imitation learning 방법보다 20.4% 더 높은 성능을 보인다. 또한, OpenVLA가 Low-Rank Adaptation 방법을 통해 사용자 GPU에서 Finetuning 될 수 있고, 다운스트림 성공률에 영향을 미치지 않고 Quantization을 통해 효율적으로 제공될 수 있다. 마지막으로, 대규모로 VLA를 학습할 수 있는 PyTorch 코드를 공개한다. (https://github.com/openvla/openvla)

Diffusion Policy:
Action 데이터를 점진적으로 노이즈화(perturb)하여 분산된 데이터를 생성. 노이즈를 제거하며 점진적으로 원래의 Action을 복원. 복원 과정에서 학습된 Action Policy을 기반으로 적절한 Action을 생성.

[ Introduction ]

Abstract에서 언급한 두 가지 문제점(폐쇄성, 적응의 어려움)이 기존 VLA 연구들의 한계이다. 그래서 본 논문은 이를 해결하깅 위해 OpenVLA라는 7B-Parameter의 오픈소스 VLA를 제안한다. 위에서 언급했듯이, Finetuning, Low-Rank Adaptation, Quantization을 통해 효율적으로 모델이 Consumer에게 제공될 수 있다. 마지막으로, 대규모로 VLA를 학습할 수 있는 PyTorch 코드를 공개한다.

[ TheOpenVLAModel ]

(1) Preliminaries:Vision-LanguageModels

- OpenVLA Architecture

위 그림처럼 대부분의 VLM의 아키텍처는 다음의 세 가지 구성 요소로 이루어져 있다.

• Visual encoder
  이미지 입력을 이미지 패치 임베딩(Image Patch Embeddings으로 매핑한다.

• Projector
  Vision Encoder의 출력 임베딩을 받아, 이를 LM의 입력 공간으로 매핑.

• LLM
  모델의 핵심 Backbone 역할을 수행.

VLM 학습 동안, 모델은 다양한 인터넷 출처에서 수집된 VLM 데이터를 사용하여 Next Text Token Prediction Objective를 통해 엔드-투-엔드 학습된다. OpenVLA의 VLM 아키텍처 OpenVLA는 Prismatic-7B VLM (600M-Parameter의 Vision Encoder-2개의 MLP Projector-7B-Parameter의 Llama 2) 을 기반으로 구축되었다.

Prismatic의 주요 특징은 CLIP 또는 SigLIP 단일 인코더와 비교할 때, DinoV2 특징을 추가로 사용하는 것이 공간적 추론(spatial reasoning)을 향상시키는 데 유용하다는 점이다. 이는 로봇 제어에서 특히 중요한 역할을 할 수 있다.

(2) OpenVLA Training Procedure

OpenVLA를 학습하기 위해, Prismatic-7B VLM Backbone을 로봇 Action Prediction Task를 위해 Fine-tuning 한다. 이때, 로봇 Action Prediction Task를 Vision-Language Task로 정의하여, 입력 관찰 이미지와 자연어 Task Instruction을 로봇 Action으로 맵핑한다.

LLM의 Backbone이 로봇 Action을 예측할 수 있도록, Continuous Robot Actions을 LLM의 토크나이저가 사용하는 Discrete Tokens으로 변환한다. 각 로봇 Action 차원을 각각 256개의 구간으로 이산화한다. 이때, 각 차원의 구간 폭은 Train 데이터의 ACtion 범위 중 1~99번째 백분위수를 균등하게 나누도록 설정한다. 이 방식으로 N차원 로봇 Action에 대해 0~255 범위의 이산 정수 N개를 얻는다.

그러나 OpenVLA에서 사용하는 Llama 토크나이저는 Finetuning 중 새롭게 추가된 토큰을 위한 100개의 특수 토큰만 받으며, 이는 256개의 Action 토큰을 담기에는 부족하다. 그래서 Llama 토크나이저에서 가장 적게 사용되는 어휘 토큰에 256개 토큰을 덮어씌워 Action 토큰으로 대체하였다.

Action이 토큰 시퀀스로 처리된 후, OpenVLA는 Next-Token Prediction Objective를 사용해 학습된다. 이때, cross-entropy loss은 Predictioned Action 토큰에서만 평가된다.

• 백분위수 사용의 이점: 데이터의 이상치가 이산화 간격을 지나치게 확장시키는 문제를 방지.

(3) Training Data

OpenVLA 학습 데이터셋을 구축하는 목표는 다양한 로봇, 장면, Tasks를 포함하여 최종 모델이 다양한 로봇을 바로 제어할 수 있도록 하는 것이다. 이를 통해 새로운 로봇 환경에 효율적으로 Finetuning할 수 있는 기반을 마련하는 것이다.

OpenX-Embodiment 데이터셋을 기반으로 학습 데이터셋을 구성한다. OpenX 데이터셋은 작성 당시 70개 이상의 개별 로봇 데이터셋과 200만 개 이상의 로봇 Trajectories으로 구성되어 있으며, 이후 대규모 커뮤니티 노력으로, 일관되고 사용하기 쉬운 데이터 형식으로 통합되었다.

학습 데이터로 사용 가능한 데이터를 실질적으로 만들기 위해, Raw dataset에 대해 여러 단계의 선별 작업을 수행했다.

• 1. 일관된 입력 및 출력 공간 보장
     Train 데이터셋을 3인칭 카메라 시점이 포함된 Manipulation datasets과 End-effector 제어로 제한하였다.
• 2. 균형 잡힌 데이터 혼합
     Data mixture weights를 활용해 첫 번째 필터링 단계를 통과한 모든 데이터셋을 처리하였다. Task 및 장면의 다양성이 낮은 데이터셋을 가중치를 줄이거나 제거하고, 더 높은 다양성을 가진 데이터셋에 가중치를 부여한다.

그 결과 DROID 데이터셋의 Action Token 정확도가 학습 내내 낮았다. 이는 이 데이터셋의 다양성을 수용하려면 더 큰 모델이나 높은 혼합 가중치가 필요하다는 점을 시사한다. 그래서 최종 모델의 성능에 영향을 주지 않기 위해, 학습의 마지막 1/3 구간에서는 DROID 데이터를 데이터 혼합에서 제외하였다.

(4) Open VLA Design Decisions

OpenVLA은 VLM 백본, 이미지 해상도, Vision-Encoder 미세 조정, Train Epoch 수, LR 등에서 다양한 실험을 통해 최적화되었다. (자세한 내용은 논문 3.4에 나와있다.)

(5) Infrastructure for Training and Inference

OpenVLA는 64개의 NVIDIA A100 GPU에서 14일간 학습하였다. 로봇 제어 시 높은 성능의 로컬 컴퓨팅 장치에 의존하지 않도록, GPU 클러스터와 원격 Inference 서버를 활용하여 성능과 접근성을 모두 고려한 설계를 완성했다. (자세한 내용은 논문 3.5에 나와있다.)

[ The OpenVLA Codebase ]

Pytorch 기반 모듈형 코드 베이스를 공개하였다. 단일 GPU에서의 Finetuning부터 대규모 클러스터 학습까지 지원한다 또한, AMP(자동 혼합 정밀도), FlashAttention, FSDP 등의 라이브러리를 사용하며, HuggingFace와도 통합하였다. 그리고 LoRA를 활용한 경량 학습 및 양자화 추론으로 메모리 최적화가 가능하게끔 하였다.

[ Experiments ]

OpenVLA의 실험 평가 목표는 다음 세 가지이다.

• (1) 비교 성능: OpenVLA가 여러 로봇 플랫폼과 일반화 작업에서 기존의 모델들과 비교하여 얼마나 뛰어난가?
• (2) Finetuning 성능: OpenVLA가 새로운 로봇 환경과 작업에 얼마나 효과적으로 Finetuning될 수 있으며, 데이터 효율적인 Imitation 학습 방법들과 비교하면 어떠한가?
• (3) 컴퓨팅 요구 감소: OpenVLA에서 LoRA를 적용한 Finetuning 및 양자화를 통해 학습 및 추론의 계산 요구를 얼마나 줄일 수 있는가? 성능과 계산 요구 간의 Trade-Offs는 무엇인가?

(1) Direct Evaluations on Multiple Robot Platforms

- BridgeData V2 WidowX robot evaluation tasks and results

위 그래프는 OpenVLA, RT-2-X, RT-1-X, Octo에 대해서 다양한 일반화에서 비교한 결과를 보여준다. OpenVLA는 평균 성능에서 70.6%로 가장 높았고, 세부적인 Task에서는 Vision, Physucak 일반화와 Language에서 가장 높은 성능을 보였다. RT-2-X는 Semantic 일반화(새로운 개념과 지시사항 이해)에서 다소 강점을 보였으나, OpenVLA와 비슷한 수준이다.

(2) Data-Efficient Adaptation to New Robot Setups

- Adapting to new robot setups

위 그래프는 OpenVLA와 다른 모델들이 Franka-Tabletop 및 Franka-DROID 환경에서 수행한 Task 성공률을 보여준다. OpenVLA는 다양한 작업 환경과 지시사항에 대해 범용적이고 안정적인 성능을 발휘했다.

평균 성능은 OpenVLA가 모든 작업에서 가장 높은 평균 성공률을 보인다는 것을 알 수 있다.

또한, Narrow Single-Instruction Tasks (단일 지시 작업, 예: "당근을 그릇에 넣기", "옥수수를 냄비에 붓기") 에서는 Diffusion Policy가 93.3% 성공률로 뛰어난 정밀성과 부드러운 동작을 보여줬다. OpenVLA는 다소 낮지만 꽤 준수한 성능을 낸다는 것을 알 수 있다.

Diverse Multi-Instruction Tasks (다양한 다중 지시 작업, 예: "접시 위에 물체 놓기", "물체를 넘어뜨리기") 에서는 OpenVLA가 다중 객체 및 언어 기반 작업에서 가장 우수하였다.

Visual Robustness (시각적 견고성, 예: "테이블 닦기") 에서는 OpenVLA가 가장 높은 시각적 견고성을 보인다는 것을 알 수 있다.

(3) Parameter-Efficient Fine-Tuning

위 표를 보면 알 수 있듯이, LoRA는 Full Fine-Tuning과 성능이 비슷하다. 하지만 8배 적은 계산 자원 필요한 이점이 있다. Last Layer Only & Frozen Vision은 자원 사용은 적지만 성능이 낮아 실용적이지 않다. Sandwich Fine-Tuning은 성능과 계산 자원의 균형이 괜찮지만, LoRA보다 효율성 낮다.

즉, LoRA는 성능과 계산 비용 간의 최적 균형을 제공하며, OpenVLA의 새로운 작업에 대한 효율적이고 실용적인 Finetuning 방법으로 추천한다.

(4) Memory-Efficient Inference via Quantization

- OpenVLA inference speed for various GPU

위 그래프에서 알 수 있듯이, 4비트 양자화는 메모리 효율성과 성능 사이의 최적 Trade-off를 제공한다. 이를 통해 OpenVLA는 소비자 GPU(예: RTX 4090)뿐 아니라, 구형 GPU(A5000)에서도 실질적으로 실행 가능하다.

[ Discussion and Limitations ]

OpenVLA는 범용 VLA 모델로, 다양한 로봇 플랫폼에서 강력한 성능 제공할 수 있다. Parameter를 효율적 Finetuning을 통해 새로운 로봇 환경에 쉽게 적응 가능하다. 하지만, 단일 이미지만 지원, 다중 이미지 및 관찰 이력 통합 필요하다. 또한, 테스트된 작업에서 90% 이상의 높은 신뢰도를 아직 제공하지 못한다.