[실습] OpenVLA로 VLA 직접 만져보기

들어가며

5편의 OpenVLA는 코드·weight·학습 데이터 모두 공개되어 있습니다. 24GB GPU 한 장(예: RTX 4090, A6000)만 있으면 추론과 LoRA finetune이 모두 가능합니다.

이번 글에서 직접 굴려볼 두 가지는 다음과 같습니다.

  1. OpenVLA 추론 — SimplerEnv 시뮬레이터에서 사전학습 모델로 task 수행
  2. LoRA finetune — 자기 task에 맞춰 OpenVLA를 가볍게 적응

GPU가 24GB 미만이면 추론은 CPU offloading으로 가능하나 매우 느립니다. finetune은 24GB 이상이 사실상 필수입니다.

0. 환경 준비

# OpenVLA 저장소
git clone https://github.com/openvla/openvla.git
cd openvla
pip install -e .

# SimplerEnv (Octo·OpenVLA 평가용 시뮬레이터)
git clone https://github.com/simpler-env/SimplerEnv --recurse-submodules
cd SimplerEnv
pip install -e ./ManiSkill2_real2sim
pip install -e .

CUDA 12.x + PyTorch 2.2 이상 환경을 권장합니다.

설치 확인:

import torch
import simpler_env
print(torch.cuda.is_available())  # True
print(simpler_env.ENVIRONMENTS)   # 사용 가능한 환경 목록

1. OpenVLA 추론 — SimplerEnv에서 코크 캔 집기

import torch
import simpler_env
from transformers import AutoModelForVision2Seq, AutoProcessor

device = "cuda"

# OpenVLA 모델 로드 (~16GB VRAM)
processor = AutoProcessor.from_pretrained("openvla/openvla-7b", trust_remote_code=True)
model = AutoModelForVision2Seq.from_pretrained(
    "openvla/openvla-7b",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).to(device)

# 시뮬레이션 환경
env = simpler_env.make("google_robot_pick_coke_can")
obs, _ = env.reset()
instruction = env.get_language_instruction()
print(f"task: {instruction}")  # "pick up the coke can"

import imageio
done = False
frames = []
while not done:
    image = obs["image"]  # (H, W, 3) RGB
    prompt = f"In: What action should the robot take to {instruction}? Out:"
    inputs = processor(prompt, image).to(device, dtype=torch.bfloat16)

    # OpenVLA의 predict_action: 7-DoF 행동을 정규화 해제된 형태로 반환
    action = model.predict_action(**inputs, unnorm_key="bridge_orig", do_sample=False)

    obs, reward, terminated, truncated, info = env.step(action)
    done = terminated or truncated
    frames.append(image)

print(f"success: {info.get('success', False)}")
imageio.mimsave("openvla_coke.mp4", frames, fps=10)

결과 해석

영상에는 로봇 팔이 코크 캔을 인식하고 잡으러 가는 모습이 나옵니다. 첫 시도에 100% 성공하지는 않지만, OpenVLA 논문의 Google robot 시뮬레이션 결과는 평균 60~70% 성공률 수준입니다.

unnorm_key="bridge_orig"는 OpenVLA가 행동을 정규화한 통계를 어디서 가져올지 지정합니다. SimplerEnv의 Google robot 환경은 bridge_orig 통계를 사용합니다. 다른 환경(WidowX 등)에서는 다른 키를 써야 합니다.

학습된 task가 아닌 instruction(예: "wave at me")을 주면 모델이 어색한 행동을 출력합니다. RT-2와 OpenVLA의 emergent capability를 직접 테스트해 보고 싶다면 instruction을 다양하게 바꿔보시기 바랍니다.

2. LoRA Finetune — 자기 task에 적응시키기

OpenVLA를 자기 데이터로 finetune하려면 LoRA가 사실상 유일한 현실 옵션입니다. 7B 전체 finetune은 80GB H100 여러 장이 필요하지만, LoRA는 24GB 한 장으로 충분합니다.

2-1. 데이터 준비 — RLDS 형식

OpenVLA는 Open X-Embodiment와 같은 RLDS 형식을 받습니다. 가장 빠른 시작은 LIBERO 데이터셋입니다.

# LIBERO 데이터셋 (Open-X 미니어처)
huggingface-cli download openvla/modified_libero_rlds --repo-type dataset

DROID 데이터셋도 같은 RLDS 형식이라 그대로 사용 가능합니다.

2-2. LoRA 설정과 학습

OpenVLA 저장소의 finetune 스크립트를 사용합니다.

torchrun --standalone --nnodes 1 --nproc-per-node 1 \
  vla-scripts/finetune.py \
  --vla_path "openvla/openvla-7b" \
  --data_root_dir <path-to-libero> \
  --dataset_name libero_spatial_no_noops \
  --run_root_dir <path-to-save> \
  --adapter_tmp_dir <tmp-dir> \
  --lora_rank 32 \
  --batch_size 8 \
  --grad_accumulation_steps 1 \
  --learning_rate 5e-4 \
  --image_aug True \
  --wandb_project openvla-libero

핵심 인자:

  • --lora_rank 32 — LoRA의 rank. 32면 보통 충분합니다.
  • --batch_size 8 — 24GB GPU 기준. 16GB라면 4로 낮춥니다.
  • --learning_rate 5e-4 — LoRA finetune의 표준값.

학습 시간은 보통 5만~10만 step. RTX 4090 기준 약 24시간입니다.

2-3. Finetuned 모델로 추론

학습이 끝나면 <run_root_dir>에 체크포인트가 저장됩니다. 추론 코드는 1번과 같지만 모델 경로만 로컬로 바꿉니다.

model = AutoModelForVision2Seq.from_pretrained(
    "<run_root_dir>/checkpoint-50000",
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
).to(device)

결과 해석

LIBERO-Spatial(공간 task)에서 finetune된 OpenVLA는 보통 80~90% 성공률을 달성합니다. baseline OpenVLA가 같은 task에서 60% 정도이므로, LoRA만으로도 큰 개선입니다.

LoRA로 학습된 가중치는 보통 100~200MB입니다. 본체 14GB와는 별도로 저장되어, 여러 task에 대한 작은 LoRA 어댑터를 swap하는 식으로 운영할 수 있습니다.

정리

실습손에 잡힌 개념
OpenVLA + SimplerEnv 추론7B VLA가 자연어 명령으로 로봇을 움직이는 모습
LoRA finetune24GB GPU 한 장으로 자기 task에 적응시키기

OpenVLA는 RT-2의 청사진을 오픈으로 복원한 모델입니다. 이 코드를 직접 만져본 뒤에는 6편의 π₀ 코드도 매우 친숙하게 읽힙니다.

다음 글 안내

LinkedIn