[논문 리뷰] Keypoint Action Tokens Enable In-Context Imitation Learning in Robotics

Norman Di Palo, Edward Johns|arXiv (Cornell University)|2024. 03. 28.

Robot Manipulation and Learning인용 수 4

한 줄 요약

이 논문은 GPT-4 Turbo와 같은 사전 훈련된 텍스트 기반 트랜스포머 모델이 시각적 관측치와 동작 궤적을 키포인트 기반 텍스트 토큰 시퀀스로 변환함으로써, 로봇 공학 분야에서 적자원(10개 이하의 예시)으로도 문맥 기반 암시적 학습을 수행할 수 있도록 하는 Keypoint Action Tokens (KAT)를 소개한다. 로봇 데이터에 대한 미세튜닝 없이도 KAT는 실제 작업에서 최신 기술 성능을 달성하며, 이는 대규모 언어 모델이 패턴 학습만으로도 시각-행동 순차 처리 작업으로 일반화할 수 있음을 보여준다.

ABSTRACT

We show that off-the-shelf text-based Transformers, with no additional training, can perform few-shot in-context visual imitation learning, mapping visual observations to action sequences that emulate the demonstrator's behaviour. We achieve this by transforming visual observations (inputs) and trajectories of actions (outputs) into sequences of tokens that a text-pretrained Transformer (GPT-4 Turbo) can ingest and generate, via a framework we call Keypoint Action Tokens (KAT). Despite being trained only on language, we show that these Transformers excel at translating tokenised visual keypoint observations into action trajectories, performing on par or better than state-of-the-art imitation learning (diffusion policies) in the low-data regime on a suite of real-world, everyday tasks. Rather than operating in the language domain as is typical, KAT leverages text-based Transformers to operate in the vision and action domains to learn general patterns in demonstration data for highly efficient imitation learning, indicating promising new avenues for repurposing natural language models for embodied tasks. Videos are available at https://www.robot-learning.uk/keypoint-action-tokens.

연구 동기 및 목표

• 로봇 공학 분야에서 로봇 데이터에 대한 미세튜닝 없이도 사전 훈련된 언어 모델을 사용해 적자원 암시적 학습을 가능하게 하기 위해.
• 대규모 텍스트 기반 트랜스포머의 잠재적 적자원 패턴 인식 능력을 활용해 로봇 학습에서의 데이터 부족 문제를 해결하기 위해.
• 시각적 관측치와 동작 시퀀스를 텍스트 호환 토큰 형식으로 매핑하는 일반 목적의 프레임워크를 개발하기 위해.
• 사전 훈련된 LLM이 낮은 데이터 환경에서 특화된 암시적 학습 모델(예: 디퓨전 정책)과 동일하거나 더 나은 성능을 내는지 평가하기 위해.
• 다양한 실제 작업, 특히 공간 일반화, 새로운 물체 일반화, 6-DoF 궤적 제어 등에 대해 KAT의 확장성과 성능을 조사하기 위해.

제안 방법

• 키포인트 동작 토큰(KAT)은 시각적 관측치와 동작 궤적을 키포인트 좌표의 시퀀스로 표현하며, 이를 자연어 토큰(예: '왼손이 (x,y,z) 위치에 있음')으로 변환해 텍스트 기반 트랜스포머에 입력한다.
• 이 방법은 각 관측 프레임에서 고정된 수의 키포인트(K)를 추출하기 위해 사전 훈련된 비전 트랜스포머를 사용하며, 공간적 및 시간적 구조를 기술적 토큰의 시퀀스로 인코딩한다.
• 예시와 현재 관측치가 프롬프트에 연결되며, 입력 시퀀스에는 전문가 예시에서 유래한 키포인트 토큰과 동작 토큰이 포함된다.
• 텍스트 사전 훈련된 트랜스포머(예: GPT-4 Turbo)는 프롬프트를 기반으로 다음 동작 토큰 시퀀스를 생성함으로써 문맥 기반 학습을 수행하며, 이 토큰들은 이후 로봇 제어 명령어로 복호화된다.
• 이 프레임워크는 완전히 추론 모드에서 작동하며, 로봇 데이터에 대한 미세튜닝이나 추가 훈련이 필요하지 않다.
• 이 방법은 장문의 입력을 지원하여 하나의 프롬프트 안에서 다수의 예시와 복잡한 궤적을 처리할 수 있다.

Figure 1: An illustration of our pipeline. KAT transforms a visual observation into a sequence of keypoint tokens . From it, the text-pretrained Transformer (LLM) predicts action tokens , that are then executed as a trajectory of poses.

실험 결과

연구 질문

• RQ1 사전 훈련된 텍스트 기반 트랜스포머가 로봇 분야에서 로봇 데이터에 대한 미세튜닝 없이도 적자원 문맥 기반 암시적 학습을 수행할 수 있는가?
• RQ2 10개의 예시만으로 KAT가 새로운 물체, 공간 구성, 다중 모odal 행동에 대해 얼마나 잘 일반화되는가?
• RQ3 로봇 전용 훈련 없이도 더 최신 세대의 LLM으로 갈수록 KAT의 성능이 향상되는가?
• RQ4 낮은 데이터 환경에서 KAT가 최신 기술 수준의 디퓨전 정책과 동일하거나 더 나은 성능을 낼 수 있는가?
• RQ5 더 긴 예시 시퀀스나 실시간 온라인 제어에 적용했을 때 KAT의 한계는 무엇인가?

주요 결과

• KAT는 10개의 예시와 미세튜닝 없이도 실제 로봇 작업인 Align T, Wipe Plate, Sweep, Bottle Upright에서 최신 기술 성능을 달성한다.
• GPT-4 Turbo는 GPT-3 Turbo와 GPT-2보다 적자원 암시적 학습에서 더 뛰어난 성능을 보이며, 더 최신의 LLM이 시각-행동 순차 처리에 더 잘 일반화됨을 보여준다.
• 낮은 데이터 환경에서 KAT는 디퓨전 정책과 동등하거나 더 나은 성능을 보이며, 이는 LLM을 활용한 문맥 기반 학습이 암시적 학습에 매우 효율적일 수 있음을 시사한다.
• LLM의 더 큰 컨텍스트 창이 성능 향상에 기여함을 확인하여, 더 긴 입력 시퀀스가 패턴 인식과 궤적 생성 능력을 향상시킨다는 점을 시사한다.
• 모델 재훈련이 필요 없기 때문에, 예시 후 즉각적인 기술 배포가 가능하다.
• KAT는 새로운 물체와 공간 구성에 대해 강력한 일반화 능력을 보이며, 실제 환경에서의 분포 이탈에 대비한 강건성을 보여준다.

Figure 2: Illustration of the extraction pipeline of keypoint tokens .