비전-언어-액션 모델 미세 조정
소개
비전-언어-행동(VLA) 모델은 시각적 인식, 언어 이해, 로봇 제어를 통합하여 자연어 명령을 따라 물리적 작업을 수행할 수 있는 통합 모델로, 구현형 AI 분야에서 획기적인 발전을 나타냅니다. 이 튜토리얼에서는 특정 로봇 응용 프로그램에 맞게 VLA 모델을 미세 조정하는 전체 과정을 다룹니다.
비전-언어-액션 모델이란 무엇인가요?
VLA는 컴퓨터 비전(이미지/비디오 해석), 자연어 처리(텍스트 이해 및 생성), 행동 실행(환경이나 시스템과 상호작용)을 결합합니다. 이를 통해 시각적 및 텍스트 입력을 기반으로 인지하고, 추론하며, 행동할 수 있습니다.
VLA가 실제로 어떻게 작동하는지 보고 싶다면, Hannah Fry와 함께하는 구글 딥마인드 로보틱스 연구소 투어를 확인해 보세요.
인기 있는 VLA 아키텍처에는 OpenVLA, RT-2(로보틱 트랜스포머 2), PaLM-E가 있습니다.
핵심 내용
- VLA 모델은 엔드 투 엔드 아키텍처에서 카메라 이미지와 자연어 지시를 로봇 동작으로 직접 매핑함으로써 시각, 언어, 행동을 통합합니다.
- LoRA를 사용하면 전체 모델 대신 0.1-0.2%의 파라미터만 학습하여 효율적인 미세 조정이 가능하며, 24GB VRAM을 가진 단일 GPU에서도 학습할 수 있습니다.
전제 조건
하드웨어 요구 사항
- 최소 24GB VRAM을 가진 GPU (RTX 3090/4090, A100, 또는 H100)
- 로봇 하드웨어 또는 시뮬레이션 환경
- 추론 실행 시 시각 입력을 위한 카메라
- 로봇 하드웨어 또는 시뮬레이션 환경
단계 0: GPU와 Jupyter Notebook 설정
Linux-Console.net GPU Droplet를 설정한 후 터미널에서 SSHroot@[IPV4]로 접속하세요
터미널에서:
# Core dependencies
pip install torch torchvision
pip install transformers accelerate
pip install datasets
pip install wandb # for experiment trackingpip install jupyter lab
jupyter lab1단계: VLA 아키텍처 이해하기
주요 구성 요소: 비전 인코더: 카메라 이미지를 시각적 임베딩으로 처리 언어 인코더: 텍스트 명령을 언어 임베딩으로 변환 융합 모듈: 시각 정보와 언어 정보를 결합 액션 디코더: 융합 표현에서 로봇 동작 예측
단계 2: 데이터셋 준비하기
데이터 수집
데이터셋에는 다음과 같은 에피소드가 포함되어야 합니다:
- 이미지: 각 시간 단계에서의 RGB 카메라 관측
- 언어: 자연어 작업 설명
- 동작: 로봇 동작 시퀀스(관절 위치, 속도 또는 엔드 이펙터 자세)
- 메타데이터: 성공 라벨, 에피소드 길이 등 등.
데이터 형식 예시
{
"episode_0": {
"images": [img_0, img_1, ..., img_T], # shape: (T, H, W, 3) where T is total frames
"language": "pick up the red block",
"actions": [act_0, act_1, ..., act_T], # shape: (T, action_dim)
"success": True
}
}데이터셋 구성
from datasets import Dataset, DatasetDict
import numpy as np
def create_vla_dataset(episodes):
"""Convert robot episodes to HuggingFace dataset format"""
data = {
"images": [],
"language": [],
"actions": [],
"episode_id": []
}
for ep_id, episode in enumerate(episodes):
for t in range(len(episode['images'])):
data['images'].append(episode['images'][t])
data['language'].append(episode['language'])
data['actions'].append(episode['actions'][t])
data['episode_id'].append(ep_id)
return Dataset.from_dict(data)
# Split into train/val
dataset = create_vla_dataset(your_episodes)
dataset = dataset.train_test_split(test_size=0.1)3단계: 기본 모델 설정
사전 학습된 VLA 로딩
from transformers import AutoModel, AutoProcessor
import torch
# Load pre-trained VLA (example using OpenVLA)
model_name = "openvla/openvla-7b"
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModel.from_pretrained(
model_name,
torch_dtype=torch.bfloat16,
device_map="auto"
)
# Freeze base model parameters (optional for LoRA)
for param in model.parameters():
param.requires_grad = False행동 공간 적응
당신 로봇의 행동 공간은 아마도 사전 학습 데이터와 다를 것입니다:
from torch import nn
class ActionHead(nn.Module):
def __init__(self, hidden_dim, action_dim):
super().__init__()
self.fc = nn.Sequential(
nn.Linear(hidden_dim, 512),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(512, action_dim)
)
def forward(self, x):
return self.fc(x)
# Add custom action head
model.action_head = ActionHead(
hidden_dim=model.config.hidden_size,
action_dim=7 # e.g., 6-DOF arm + gripper
)4단계: LoRA를 활용한 효율적인 파인튜닝
저랭크 적응(LoRA)은 학습 가능한 파라미터 수를 줄여 효율적인 미세 조정을 가능하게 합니다:
from peft import LoraConfig, get_peft_model
# Configure LoRA
lora_config = LoraConfig(
r=16, # rank
lora_alpha=32,
target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
lora_dropout=0.1,
bias="none",
task_type="CAUSAL_LM"
)
# Apply LoRA to model
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
#Expected output: trainable params: 8.3M || all params: 7B || trainable%: 0.12%5단계: 데이터 처리 및 증강
이미지 전처리
from torchvision import transforms
image_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.RandomHorizontalFlip(p=0.5),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])행동 정규화
def normalize_actions(actions, stats):
"""Normalize actions to [-1, 1] range"""
return (actions - stats['mean']) / (stats['std'] + 1e-8)
def denormalize_actions(normalized_actions, stats):
"""Convert back to original action space"""
return normalized_actions * (stats['std'] + 1e-8) + stats['mean']
# Calculate statistics from your dataset
action_stats = {
'mean': dataset['train']['actions'].mean(axis=0),
'std': dataset['train']['actions'].std(axis=0)
}데이터로더 설정
from torch.utils.data import DataLoader
def collate_fn(batch):
"""Custom collate function for VLA data"""
images = torch.stack([processor.image_processor(b['images'])
for b in batch])
text = processor.tokenizer([b['language'] for b in batch],
padding=True, return_tensors="pt")
actions = torch.tensor([b['actions'] for b in batch],
dtype=torch.float32)
return {
'pixel_values': images,
'input_ids': text['input_ids'],
'attention_mask': text['attention_mask'],
'actions': actions
}
train_loader = DataLoader(
dataset['train'],
batch_size=32,
shuffle=True,
collate_fn=collate_fn,
num_workers=4
)6단계: 학습 루프
손실 함수
import torch.nn.functional as F
def vla_loss(predicted_actions, target_actions, reduction='mean'):
"""Action prediction loss (MSE for continuous actions)"""
mse_loss = F.mse_loss(predicted_actions, target_actions, reduction=reduction)
return mse_loss
# Alternative: Action chunking for temporal coherence
def chunked_action_loss(pred_chunks, target_chunks, chunk_size=10):
"""Predict multiple future actions at once"""
loss = 0
for i in range(chunk_size):
loss += F.mse_loss(pred_chunks[:, i], target_chunks[:, i])
return loss / chunk_size훈련 스크립트
from accelerate import Accelerator
from tqdm import tqdm
import wandb
# Initialize accelerator for distributed training
accelerator = Accelerator(mixed_precision='bf16')
# Setup
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10000)
# Prepare for distributed training
model, optimizer, train_loader = accelerator.prepare(
model, optimizer, train_loader
)
# Initialize wandb
wandb.init(project="vla-finetuning", config={
"learning_rate": 1e-4,
"batch_size": 32,
"epochs": 10
})
# Training loop
global_step = 0
for epoch in range(10):
model.train()
epoch_loss = 0
for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}"):
# Forward pass
outputs = model(
pixel_values=batch['pixel_values'],
input_ids=batch['input_ids'],
attention_mask=batch['attention_mask']
)
# Predict actions
predicted_actions = model.action_head(outputs.last_hidden_state[:, -1])
# Compute loss
loss = vla_loss(predicted_actions, batch['actions'])
# Backward pass
accelerator.backward(loss)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
# Logging
epoch_loss += loss.item()
global_step += 1
if global_step % 100 == 0:
wandb.log({
"loss": loss.item(),
"learning_rate": scheduler.get_last_lr()[0],
"epoch": epoch
})
print(f"Epoch {epoch+1} - Average Loss: {epoch_loss / len(train_loader):.4f}")
# Save checkpoint
if (epoch + 1) % 2 == 0:
accelerator.save_model(model, f"checkpoints/epoch_{epoch+1}")단계 7: 평가
시뮬레이션 평가
def evaluate_in_simulation(model, env, num_episodes=50):
"""Evaluate model in simulation environment"""
model.eval()
success_count = 0
with torch.no_grad():
for episode in range(num_episodes):
obs = env.reset()
instruction = env.get_task_instruction()
done = False
while not done:
# Process observation
image = torch.tensor(obs['image']).unsqueeze(0)
text_inputs = processor.tokenizer(instruction, return_tensors="pt")
# Predict action
outputs = model(
pixel_values=image.to(model.device),
input_ids=text_inputs['input_ids'].to(model.device)
)
action = model.action_head(outputs.last_hidden_state[:, -1])
# Denormalize and execute
action = denormalize_actions(action.cpu().numpy(), action_stats)
obs, reward, done, info = env.step(action[0])
if info['success']:
success_count += 1
success_rate = success_count / num_episodes
print(f"Success Rate: {success_rate*100:.1f}%")
return success_rate행동 복제 지표
def evaluate_bc_metrics(model, val_loader):
"""Evaluate behavioural cloning performance"""
model.eval()
total_mse = 0
total_cosine_sim = 0
n_samples = 0
with torch.no_grad():
for batch in val_loader:
outputs = model(
pixel_values=batch['pixel_values'],
input_ids=batch['input_ids']
)
pred_actions = model.action_head(outputs.last_hidden_state[:, -1])
# MSE
mse = F.mse_loss(pred_actions, batch['actions'])
total_mse += mse.item() * len(batch['actions'])
# Cosine similarity
cosine_sim = F.cosine_similarity(pred_actions, batch['actions']).mean()
total_cosine_sim += cosine_sim.item() * len(batch['actions'])
n_samples += len(batch['actions'])
return {
'mse': total_mse / n_samples,
'cosine_similarity': total_cosine_sim / n_samples
}단계 8: 배포
모델 내보내기
# Save fine-tuned model
model.save_pretrained("./finetuned_vla")
processor.save_pretrained("./finetuned_vla")
# For deployment, merge LoRA weights
from peft import PeftModel
base_model = AutoModel.from_pretrained(model_name)
merged_model = PeftModel.from_pretrained(base_model, "./finetuned_vla")
merged_model = merged_model.merge_and_unload()
merged_model.save_pretrained("./deployed_model")실시간 추론
class VLAController:
def __init__(self, model_path):
self.model = AutoModel.from_pretrained(model_path)
self.processor = AutoProcessor.from_pretrained(model_path)
self.model.eval()
self.model.to('cuda')
@torch.inference_mode()
def predict_action(self, image, instruction):
"""Real-time action prediction"""
# Preprocess
inputs = self.processor(
images=image,
text=instruction,
return_tensors="pt"
).to('cuda')
# Predict
outputs = self.model(**inputs)
action = self.model.action_head(outputs.last_hidden_state[:, -1])
return action.cpu().numpy()[0]
# Usage
controller = VLAController("./deployed_model")
action = controller.predict_action(camera_image, "pick up the cup")
robot.execute_action(action)자원
오픈VLA
RT-2 논문과 RT-2 웹사이트: 이것이 VLA(비전-언어-행동 모델)라는 용어를 만든 작업입니다
PEFT 라이브러리
RoboMimic (데이터셋 처리용)
-
OpenVLA: 김무진의 LeRobot 연구 발표 #5
결론
VLA 모델을 미세 조정하면 로봇이 자연어로 제어되는 특화된 작업을 수행할 수 있습니다. 이 튜토리얼이 사전 학습된 모델을 특정 하드웨어와 작업에 맞게 조정하는 방법을 보여주고, 대규모 사전 학습의 장점을 활용하면서도 응용에 맞게 커스터마이징할 수 있기를 바랍니다. 여기서는 시뮬레이션부터 시작하고, 빠르게 반복하며, 모델이 개선됨에 따라 점차 실제 하드웨어로 전환하는 것이 좋습니다.
Linux-Console.net 커뮤니티와 함께 학습해 주셔서 감사합니다.