SeeQ: Training Generalist Value Functions for Long-Horizon Robotic Manipulation

기술 요약: SeeQ: 장기 호라이즌 로봇 조작을 위한 범용 가치 함수 학습

문제 정의

모방 학습을 통해 훈련된 범용 로봇 정책은 복잡하고 장기적인(long-horizon) 조작 작업 수행에 어려움을 겪습니다. 이러한 작업은 종종 여러 단계로 구성되거나, 성공하기 위해 특정 단계에 대한 반복적인 시도와 숙고를 필요로 합니다. 이러한 환경에서는 오류가 시간이 지남에 따라 누적되며, 전문가 데이터셋에는 회복(recovery) 행동이 충분히 나타나지 않습니다. Q-가치 함수는 후보 행동의 순위를 매김으로써 정책을 유도하는 메커니즘을 제공하지만, 이를 장기 호라이즌 작업에 적용하는 것은 매우 어렵습니다. 기존 방식들은 다음과 같은 세 가지 주요 장애물에 직면합니다:

1. 장기 신용 할당 호라이즌 (Long Credit-Assignment Horizons): (최종 성공/실패와 같은) 희소한 태스크 수준 보상으로부터 학습하는 것은 긴 시퀀스에 걸쳐 신호를 전파해야 하므로 어려운 벨만 백업(Bellman backup)을 초래합니다.
2. 데이터 커버리지: 오프라인 데이터셋은 장기 궤적 중에 발생하는 다양한 상태-행동 쌍을 충분히 포함하지 못하는 경우가 많아, 시간 차(Temporal-Difference, TD) 학습을 신뢰하기 어렵게 만듭니다.
3. 취약성 (Brittleness): 효과적인 가치 학습 없이는, 정책이 진행을 만드는 행동과 실패로 이어지는 행동을 구분할 수 없습니다. 특히 정밀도가 요구되거나 다단계인 작업에서 더욱 그러합니다.

기존 방법들은 TD 학습을 피하거나(Monte Carlo 리턴 사용, 이는 데이터 분포 이상의 정책 개선을 제한함), 전체 태스크 호라이즌에 대해 TD 학습에 의존하지만, 이는 오류 누적 문제를 겪습니다.

방법론: SeeQ (Subtask-elicited Q-functions)

저자들은 전체 태스크가 아닌 **현재 활성화된 서브태스크(subtask)**에 집중함으로써 학습 호라이즌을 단축하는 범용 Q-가치 함수 훈련 프레임워크인 SeeQ를 제안합니다.

핵심 아키텍처 및 훈련

SeeQ는 사전 훈련된 시각-언어 모델(VLM) 백본(구체적으로 3B 규모의 PaliGemma 모델)을 활용하며, 두 단계의 훈련 과정을 도입합니다:

1. 범용 사전 훈련 (Generalist Pretraining): 모델은 서브태스크 주석이 포함된 다양한 오픈 소스 로봇 조작 데이터셋(예: RoboCOIN)을 사용하여 사전 훈련됩니다. 이 단계는 다양한 형태(embodiment)에 걸쳐 태스크 진행 상황과 행동 조건화에 대한 모델의 이해를 정규화합니다.
2. 다운스트림 미세 조정 (Downstream Finetuning): 모델은 특정 타겟 작업에 맞춰 미세 조정됩니다.

주요 기술 혁신

• 서브태스크 수준의 가치 예측: Q-함수는 최종 태스크 완료에 대한 리턴을 예측하는 대신, 활성화된 서브태스크($\tilde{l}_t$)에 대한 기대 리턴을 추정합니다. 이는 예측 호라이즌을 줄여 TD 학습을 더 안정적이고 효과적으로 만듭니다.
• 자기회귀적 서브태스크 디코딩 (Autoregressive Subtask Decoding): 테스트 시점에 인간의 주석이나 외부 서브태스크 예측기가 필요하지 않도록, Q-함수 아키텍처를 현재 상태와 태스크 명령어를 바탕으로 자연어로 활성 서브태스크를 자기회귀적으로 예측하도록 수정했습니다.
  • 훈련 중에는 실제 서브태크 주석에 대한 다음 토큰 예측 손실(next-token prediction loss)로 감독됩니다.
  • 추론 시에는 모델이 현재 상태와 태스크 명령을 기반으로 서브태스크 텍스트($\hat{l}_t$)를 생성한 후, 이 생성된 텍스트에 조건화하여 가치를 예측합니다.
• TD-BoN (Temporal-Difference with Best-of-N) 학습: 훈련 목적 함수는 서브태스크 수준의 TD 학습과 "Best-of-N" 백업 전략을 결리합니다.
  • 후보 액션 청크(action chunks)는 베이스 정책($\pi_{base}$)으로부터 샘플링됩니다.
  • 타겟 가치는 $N$개의 후보 중 가장 높은 추정 가치를 가진 액션 청크를 사용하여 계산됩니다.
  • 이 접근 방식은 훈련 백업을 테스트 시점의 스티어링 절차와 일치시키며, Q-함수가 데이터셋에서 관찰된 것보다 더 나은 행동을 평가할 수 있게 합니다.
• 경계 간 부트스트래핑 없음 (No Bootstrapping Across Boundaries): TD 타겟은 서브태스크 경계를 넘어 부트스트래핑되지 않습니다. 액션 청크 호라이즌 내에서 서브태스크가 종료되는 경우, 백업 항은 0이 되어 가치 함수가 현재 서브태스크의 진행 상황만을 엄격하게 반영하도록 보장합니다.

전체 손실 함수는 서브태스크 가치에 대한 TD 손실과 서브태스크 텍스트에 대한 다음 토큰 예측 손실을 결합합니다:
$$\mathcal{L}_{SeeQ}(\theta) = \mathcal{L}_{TD}(\theta) + \lambda_{subtask}\mathcal{L}_{subtask}(\theta)$$

테스트 시점 추론

배포 시, SeeQ는 Best-of-N 선택을 통해 베이스 정책($\pi_{base}$)을 유도(steer)합니다:

1. 상태 $s_t$와 태스크 명령어 $l$이 주어지면, 모델은 활성 서브태스크 $\hat{l}_t$를 자기회귀적으로 예측합니다.
2. 베이스 정책은 $N$개의 후보 액션 청크를 제안합니다.
3. Q-함수는 $s_t$, $l$, 그리고 예측된 $\hat{l}_t$에 조건화되어 각 청크의 점수를 매깁니다.
4. 가장 높은 점수를 받은 액션 청크가 실행됩니다.

주요 기여

1. 서브태스크 수준의 정식화: 본 논문은 장기 호라이즌 가치 학습을 단기 호라이즌인 서브태스크 수준 학습으로 재구성하는 방법을 도입하여, 장기 호라이즌에서의 희소한 보상 문제를 효과적으로 우회합니다.
2. 언어 조건부 서브태스크 추론: 자연어로 활성 서브태스크를 명시적으로 예측하는 새로운 아키텍처를 통해, 테스트 시점에 모듈형 서브태스크 예측 시스템이나 인간의 주석이 필요 없도록 했습니다.
3. 범용 사전 훈련 전략: 다양한 로봇 데이터로 VLM 백본을 사전 훈련하는 것이 강건한 액션 조건화 학습과 다운스트림 미세 조정 시 특정 이미지 특징에 대한 과적합을 줄이는 데 필수적임을 입증했습니다.
4. 실험적 검증: 두 가지 로봇 플랫폼에서 네 가지 실제 양손 조작 작업(셔츠 걸기, 뚜껑 닫기, 식료품 포장, 레고 분해)에 대해 광범위한 평가를 수행했습니다.

결과

저자들은 변형 가능한 물체, 정밀한 정렬, 다단계 협업이 포함된 네 가지 실제 작업에 대해 SeeQ를 평가했습니다.

• 정책 스티어링 성능: SeeQ는 모든 작업에서 베이스 정책의 성공률을 실질적으로 향상시켰습니다.
  • 셔츠 걸기 (Shirt-hang): 10/24 (41.7%)에서 22/24 (91.7%)로 향상.
  • 뚜껑 닫기 (Lid-sealing): 10/24 (41.7%)에서 15/24 (62.5%)로 향상.
  • 식료품 포장 (Grocery-packing): 9/24 (37.5%)에서 17/24 (70.8%)로 향상.
  • 레고 분해 (Lego-disassembly): 5/24 (20.8%)에서 10/24 (41.7%)로 향상.
  • 전체적으로 평균 성공률이 35.4%에서 66.7%로 증가했습니다 (1.88배 개선).
• 절제 연구 (Ablation Studies):
  • 사전 훈련: 로봇 데이터 사전 훈련을 제거하면 성능이 크게 저하되었습니다 (셔츠 걸기 기준 22/24에서 8/24로). 이는 일반화를 위한 다양한 데이터의 필요성을 강조합니다.
  • 서브태스크 조건화: 서브태스크를 명시적으로 디코딩하고 이에 조건화하는 것이 매우 중요했습니다. 서브태스크 예측을 제거하면 베이스 정책보다 낮은 성능(8/24)을 보였으며, 예측은 추가하되 조건화는 하지 않으면 15/24로 개선되었습니다. 전체 SeeQ는 22/24를 달려냈습니다.
  • 베이스라인과의 비교: SeeQ는 태스크 수준의 Monte Carlo 회귀, 태스크 수준의 TD 학습, 그리고 (Best-of-N 샘플링 대신 데이터셋 액션을 사용하는) 서브태스크 수준의 SARSA보다 우수한 성능을 보였습니다.

의의 및 주장

본 논문은 범용 Q-함수를 훈련할 때 서브태스크 수준의 가치가 희소한 장기 호라이즌 성공 신호보다 더 효과적인 학습 목표를 제공한다고 주장합니다. 조작 작업을 중간 마일스톤으로 분해하는 자연스러운 특성을 활용함으로써, SeeQ는 장기 호라이즌에서 흔히 발생하는 오류 누적 없이 효과적인 TD 학습을 가능하게 합니다.

저자들은 이 접근 방식이 최소한의 미세 조정을 통해 새로운 작업에 적용될 수 있는 범용 가치 학습을 가능하게 한다고 강조합니다 (단, 해당 작업이 서브태스크로 분해 가능하다는 전제 하에). 서브태스크를 예측하기 위해 언어를 명시적으로 사용하는 것은 가치 추정을 VLM의 텍스트 생성 능력과 일치시켜, 서브태스크 전환 시의 강건성을 높입니다.

이 연구는 단기 호라이즌 학습 목표와 언어 구조적 추론을 결면하는 것이, 대규모 로봇 데이터 사전 훈련을 활용할 때 더욱 강건한 장기 호라이즌 로봇 조작을 향한 유망한 경로임을 시사합니다. 저자들은 서브태스크 경계의 모호함과 베이스 정책의 후보 액션 품질에 대한 의존성과 같은 한계를 인정하면서도, 이 프레임워크가 복잡한 실제 로보틱스를 위한 가치 함수를 실용화하는 데 있어 중요한 진전임을 밝히고 있습니다.