기술 요약: RoboSSM – 상태 공간 모델을 통한 확장 가능한 인컨텍스트 모방 학습

1. 문제 정의

인컨텍스트 모방 학습(In-context Imitation Learning, ICIL)은 로봇이 배포 시 파라미터 업데이트 없이 소수의 시연(demonstrations)을 통해 새로운 작업에 적응할 수 있게 한다. 최근의 ICIL 방법론들은 트랜스포머(Transformer) 기반 아키텍처를 채택하여 모방 학습을 시퀀스 모델링 문제로 정식화함으로써 성공적인 성과를 거두었으나, 심각한 확장성 한계에 직면해 있다. 트랜스포머는 시퀀스 길이($L$)에 대해 이차 시간 복잡도($O(L^2)$)를 보이며, 훈련 시 보았던 프롬프트 길이를 넘어 외삽(extrapolate)하는 데 어려움을 겪는다. 결과적으로 기존의 ICIL 방법론들은 훈련 시 사용된 것보다 더 긴 테스트 타임 프롬프트(예: 더 많은 시연 또는 긴 호라이즌 작업)를 처리할 때 성능이 급격히 저하되며, 이는 광범ful한 컨텍스트를 요구하는 복잡한 실제 환경 적용을 제한한다.

2. 방법론: RoboSSM

저자들은 트랜스포머를 대신하여 상태 공간 모델(State-Space Models, SSMs), 구체적으로는 Longhorn을 백본 아키텍처로 사용하는 확장 가능한 ICIL 프레임링인 RoboSSM을 제안한다.

아키텍처 및 인코딩

• 입력 처리: RoboSSM은 다중 모달 관측치(CNN을 통한 시각 데이터 및 MLP를 통한 고유 수용성 데이터)를 임베딩으로 인코딩한다. 특히, 모델이 오직 인컨텍스트 시연으로부터 작업 의도를 추론하도록 강제하기 위해 입력 표현에서 액션 토큰과 작업 언어 지시문을 제외한다.
• Longhorn 백본: 관측 임베딩의 시퀀스는 선형 재귀(linear recurrence)를 통해 업데이트되는 메모리 상태 행렬 $s_t$를 유지하는 Longhorn에 의해 처리된다. 트랜스포머의 셀프 어텐션 메커니즘과 달리, Long-horn은 온라인 볼록 최적화(online convex programming)에서 유도된 재귀적 업데이트 규칙을 사용한다:
  $$s_t = \text{argmin}_{s} \left( \|s - s_{t-1}\|_F^2 + \|s k_t - x_t\|_{\text{diag}(\beta_t)}^2 \right)$$
  이 공식은 이전 상태 정보를 유지하는 것과 새로운 입력을 통합하는 것 사이의 균형을 맞춘다.
• $\beta$-스케일링: ICIL 설정에서 데모 프롬프트에 대한 주의(attention)를 특별히 촉진하기 위해, 저자들은 가중치 벡터 $\beta_t$가 $\gamma \in (0, 1]$인 인자로 스케일링되는 테스트 타임 스케일링 메커니즘을 도입했다. 이는 상태 유지와 새로운 정보 통합 사이의 트레이드오프를 조절한다.

훈련 및 추론

• 훈련: 모델은 프롬프트(한 세트의 $N_{train}$ 데모 궤적)와 쿼리 궤적으로 구성된 멀티 태스크 데이터셋에서 훈련된다. 모델은 프롬프트에 조건화되어 쿼리 궤적에 대한 액션을 예측한다.
• 추론: 테스트 타임에 정책(policy)은 $N_{test}$개의 데몬스트레이션을 포함하는 프롬프트를 사용하여 미학습 작업에 적응한다. 모델은 초기 관측치와 프롬프트에 의해 제공된 컨텍스트를 바탕으로 파라미터 업데이트 없이 반복적으로 액션을 예측한다.

3. 주요 기여

1. 확장 가능한 ICIL 프레임워크: SSM을 인컨텍스트 모방 학습에 활용한 최초의 접근법인 RoboSSM을 도입하여, 선형 시간 추론($O(L)$)과 강력한 외삽 능력을 가능하게 했다.
2. 장기 컨텍스트 외삽: RoboSSM이 훈련 시보다 최대 16배 더 긴 테스트 타임 프롬프트를 효과적으로 처리할 수 있음을 입증하였으며, 이는 트랜스포머 기반 방법론들이 실패하는 지점이다.
3. 효율적인 추론: 프롬프트 길이가 증가함에 따라 SSM의 선형 복잡도 덕분에 트랜스포머 베이스라인에 비해 현저히 낮은 추론 실행 시간을 달성했다.
4. 장기 호라이즌 작업으로의 일반화: 짧은 호라이즌 훈련 분포로부터 긴 호라이즌 및 시간 팽창(time-dilated) 시나리오로 효과적으로 일반화할 수 있음을 보여주었다.

4. 실험 결과

실험은 다양한 조작 작업(Object, Goal, Spatial, Long, 90 suite)을 포함하는 LIBERO 벤치마크에서 수행되었다.

• 프롬프트 길이 일반화:
  • 2개의 데모($N_{train}=2$)로 훈련했을 때, RoboSSM은 테스트 데모 수($N_{test}$)가 32까지 증가하더라도 성공률을 유지하거나 오히려 향상시켰다.
  • 반면, 트랜스포머 기반 베이스라인(ICRT)은 $N_{test} > N_{train}$인 경우 성능이 급격히 떨어지며 성공률이 0%로 붕괴되었다.
• 장기 호라이즌 및 시간 팽창:
  • RoboSSM은 장기 호라이즌 작업(예: "물체 A와 물체 B를 모두 물체 C에 넣으시오") 및 시간 팽창된 데모(최대 16배 시간 신축) 작업을 성공적으로 해결했으나, ICRT는 완전히 실패했다.
• 인-디스트리뷰션(In-Distribution) 성능:
  • $N_{train} = N_{test}$인 경우에도 RoboSSM은 대부분의 시나리오, 특히 LIBERO-90 Study 및 Living Room 장면에서 ICRT를 일관되게 앞질렀다.
• 멀티 태스크 학습(MTL)과의 비교:
  • RoboSSM과 ICRT는 표준 MTL 베이스라인(언어 지시가 필요하고 인컨텍스트 프롬프팅이 없는 방식)보다 유의미하게 높은 성능을 보였으며, 이를 통해 ICIL 패러다임의 효용성을 확인했다.
• 잠재 공간 분석:
  • 잠재 표현(latent representation) 시각화를 통해 RoboSSM이 프롬프트 데모 클러스터 근처에 머무는 궤적을 생성함을 보여주었으며, 이는 안정적인 외삽을 나타낸다. 반대로 ICRT의 궤적은 컨텍스트 길이가 증가함에 따라 프롬프트 매니폴드에서 크게 벗어났다.
• 런타임:
  • RoboSSM은 프롬프트 길이에 따라 거의 선형적으로 스케일링되는 반면, ICRT의 런타임은 이차적으로 증가하여 긴 컨텍스트에서 RoboSSM 쪽으로 격차가 벌어지는 성능 우위를 보였다.

5. 의의 및 주장

본 논문은 상태 공간 모델을 효율적이고 확장 가능한 백본으로 확립함으로써 인컨텍스트 모방 학습을 위한 패러다임 전환을 제시한다고 주장한다. 주요 의의는 트랜스포머의 계산적 및 외삽적 병목 현상을 극복하여 로봇이 다음과 같은 능력을 갖추게 한 데 있다:

• 재훈련 없이 테스트 타임에 훨씬 더 긴 컨텍스트 윈도우를 처리할 수 있음.
• 소수의 프롬프트를 사용하여 미학습, 장기 호라이즌 및 시간 가변적 작업에 적응할 수 있음.
• 작업별 파라미터 업데이트 없이 새로운 데모 프롬프트를 수용함으로써 평생 학습(lifelong learning)을 위한 지속적인 적응을 지원함.

저자들은 RoboSSM이 긴 컨텍스트 ICIL을 위한 SSM의 잠재력을 보여주었지만, 일반화 능력을 더욱 높이기 위해서는 향arian 더 넓고 다양한 훈련 코퍼스가 필요할 수 있다고 결론지었다. 또한, 실제 로봇 검증이 본 연구의 핵심 기여는 아니지만, 향상된 일반화 성능을 보여줄 수 있는 향후 연구 분야임을 명시하였다.