RoboAlign: 시각-언어-행동 모델의 언어-행동 정렬을 위한 테스트 시점 추론 학습

기존 로봇 AI는 언어적 이해와 실제 물리적 행동 사이의 간극으로 인해 복잡한 작업을 수행할 때 성능이 불안정했다. 이 논문은 로봇이 행동하기 전 '생각(추론)'하는 과정을 강화학습으로 직접 교정하여, 아주 적은 데이터만으로도 로봇의 작업 성공률을 획기적으로 높이는 방법론을 제시한다.

기존 VLA 모델은 MLLM 위에 행동 헤드를 얹어 학습하지만, 언어적 추론(Reasoning)과 물리적 행동(Action) 사이의 모달리티 격차로 인해 추론 능력이 실제 행동 개선으로 이어지지 않는 한계가 있었다. 특히 정답 경로만 학습하는 SFT 방식은 모델이 다양한 상황에서 유연하게 사고하도록 유도하기 어렵다.

ROBOALIGN은 로봇이 행동 토큰을 내뱉기 전 스스로 '생각'하는 과정을 거치게 하고, 이 생각이 실제 행동의 정확도로 이어지는지를 강화학습(GRPO)으로 평가한다. 즉, '컵을 집으려면 먼저 팔을 뻗어야 한다'는 추론이 실제 '팔 뻗기' 행동 토큰의 정확도를 높이는 방향으로 모델의 내부 표현(Embedding)을 정렬한다.

이를 통해 모델은 단순히 패턴을 암기하는 것이 아니라, 주어진 시각 정보와 명령어를 바탕으로 논리적인 단계를 밟아 정확한 행동을 도출하는 법을 배운다. 결과적으로 복잡하고 긴 단계가 필요한 작업(Long-horizon tasks)에서 기존 방식보다 훨씬 뛰어난 적응력을 보여주며, 모델의 내부 특징 공간이 물리적 상태를 더 잘 구분할 수 있도록 강화된다.

FAST 토큰화 및 SFT 단계가 선행된다. 로봇의 연속적인 행동(위치, 방향, 그리퍼 상태) 벡터를 입력으로 받아 이산 코사인 변환(DCT)을 수행하고, 이를 양자화하여 이산적인 토큰으로 변환한다. 이후 일반적인 시각-언어 데이터와 로봇 행동 데이터를 혼합하여 SFT를 진행함으로써 모델이 기본적으로 행동 토큰을 생성할 수 있는 능력을 갖추게 한다.

GRPO 기반 추론 정렬 단계를 수행한다. SFT된 모델을 바탕으로 ... 태그 내에 추론 과정을 생성하게 한다. GRPO 알고리즘을 사용하여 동일한 질문에 대해 여러 응답을 샘플링하고, 각 응답의 행동 토큰이 정답과 얼마나 일치하는지를 보상으로 주어 정책을 업데이트한다.

보상 함수는 형식 보상 $r_f$와 정확도 보상 $r_a$의 산술 평균으로 설계된다. 정확도 보상 $r_a$는 생성된 토큰 시퀀스 $T_{1:n}^{gen}$과 타겟 시퀀스 $T_{1:m}^{target}$ 두 값을 입력으로 받아, 앞에서부터 일치하는 토큰의 개수를 계산하고 이를 전체 길이 $m$으로 나누어 출력한다. 이 값이 1에 가까울수록 모델이 실제 로봇의 정답 동작에 더 가깝게 추론했음을 의미하며, 이 숫자를 최대화하는 방향으로 가중치를 갱신한다.

LIBERO 벤치마크에서 ROBOALIGN은 평균 성공률 86.8%를 기록하여 SFT 베이스라인(81.5%) 및 기존 모델인 ThinkAct(84.4%)를 능가했다. 특히 긴 단계가 필요한 'Long' 카테고리에서 70.0%의 성공률을 보여 다른 방식들(약 2% 개선)에 비해 압도적인 성능 향상을 보였다.

CALVIN 벤치마크에서는 연속 작업 성공률이 2.57로 향상되었으며, 실제 로봇 환경에서도 SFT 모델(32.3%) 대비 106.6% 향상된 66.7%의 성공률을 달성했다. 이는 단 12.8K개의 추가 RL 데이터를 사용하여 얻은 결과로, 데이터 효율성이 매우 높음을 입증했다.

표현 분석 결과, ROBOALIGN을 거친 모델의 내부 표현을 이용한 KNN 분류 정확도가 39.06%에서 69.79%로 급증했다. 이는 강화학습 단계가 모델로 하여금 물리적 상태를 더 명확하게 구분할 수 있는 특징을 학습하게 했음을 시사한다.

아키텍처는 Qwen2.5-VL-7B-Instruct를 백본으로 사용하며, 그 위에 확산 모델 기반의 행동 헤드(Diffusion-based action head)를 추가한 구조이다. 학습 시에는 백본의 숨겨진 상태(Hidden states)를 입력으로 받아 행동을 예측하도록 설계되었다.

FAST(Efficient Action Tokenization) 기법을 채택하여 7차원의 로봇 상태 벡터를 주파수 영역(DCT)으로 변환 후 양자화하여 토큰화한다. 이는 긴 행동 시퀀스를 짧은 토큰 열로 압축하여 MLLM의 컨텍스트 윈도우 부담을 줄이고 학습 효율을 높인다.

GRPO(Group Relative Policy Optimization)는 별도의 가치 함수 네트워크 없이 샘플 그룹 내의 상대적 보상을 사용하여 이점(Advantage)을 계산한다. 이를 통해 메모리 점유율을 낮추면서도 대규모 모델의 강화학습을 가능하게 하여 추론과 행동 간의 정렬을 수행한다.

실험 결과에 따르면, 추론 과정(Reasoning trace)을 포함하는 것이 단순 행동 예측보다 성능이 우수하며, 특히 고수준 행동 계획이나 2D 궤적 예측보다 저수준 행동 토큰과의 직접적인 정렬이 로봇 제어 성능 향상에 더 효과적임이 확인되었다.

안전하지 않은 행동 궤적 데이터로 학습될 경우 모델이 위험한 행동을 생성하도록 편향될 가능성이 있으며, 배포 후 어떤 데이터로 학습되었는지 정확히 파악하기 어렵다는 점이 한계로 언급되었다.