안녕하세요. 이번에 가져온 논문은 π 시리즈의 첫 논문인 π0를 가져왔습니다. 꾸준히 발전되어 π0.7까지 나온 것 같은데, 이런 최신 논문을 따라가기 위해 첫 논문을 자세히 읽어보자는 취지도 있었고, π 시리즈가 다른 다양한 논문들에 영향을 주면서 로보틱스 팀 내의 대화에서도 자주 언급되는데 명확하게 모르고 있는 것 같아서 자세히 읽어보고 싶었습니다. 그럼 리뷰 시작하겠습니다.

Abstract

Robot learning은 데이터 부족, generalization, robustness 문제로 인해 real world에서 범용적으로 동작하는 로봇을 만드는 데 어려움이 있습니다. 저자는 Internet-scale semantic knowledge를 활용하기 위해, 사전학습된 VLM에 flow matching 기반 action generation 구조를 결합한 pi-0를 제안합니다. pi-0는 generalist robot policy(robot foundation model)라고 부를 수 있고, 다양한 로봇에서 수집한 대규모 데이터를 학습할 수 있도록 설계했다고 합니다.

Overview

저자는 pi-0의 학습과 아키텍처에 대해서 간략하게 정리하고 넘어갑니다. 자체적으로 수집한 데이터(7개의 다양한 robot configuration, 68개의 task)와 OXE dataset을 pre-training mixture로 구성했습니다. 여기서 language label도 사용했는데, task 이름과 2초 정도 길이의 sub-trajectory에 대한 fine-grained labels도 붙여서 사용했다고 합니다. 이렇게 pre-training mixture를 구성하여 pre-training에서 언어 명령을 따르고 다양한 task를 기본적인 수준으로 수행할 수 있는 base model을 만든다고 합니다. 그 후에 복잡하고 정교한 task 수행을 위해 high-quality qurated data를 사용하여 post-training을 적용합니다.
pi-0는 PaliGemma VLM을 backbone으로 사용하고, 위에서 설명한 학습 절차를 적용합니다. PaliGemma VLM으로 선택한 이유는 간편하고 크기가 작아서 real-time에 적합해서 선택했다고 합니다. 이 구조에서 연속적인 action distribution을 생성할 수 있도록 flow matching 기반 action expert를 붙였습니다.

The Pi-0 Model

pi-0는 image와 language를 처리하는 VLM backbone과 robot state와 action을 처리하는 별도의 action expert를 결합한 구조입니다. 이 구조에서 conditional flow matching을 사용하여 continuous action distribution을 모델링합니다. 저자는 flow matching이 높은 precision과 action의 multimodal distribution 표현에 적합하다고 언급합니다. 현재 observation o_t가 주어질 때, action chunk A_t의 분포 p(A_t | o_t)를 모델링합니다.

A_t = [a_t, a_{t+1}, ... , a_{t+H-1}] (H = 50) , o_t = [I^{1}_{t}, ... , I^{n}_{t},l_t,q_t]
a 는 action, H는 chunk 길이, I는 image, n은 이미지 개수(image source: camera), l은 language, q는 joint angle입니다. 각 action chunk의 action은 action token으로 변환되어 action expert에 입력됩니다. 학습에서는 GT action chunk A_t에 Gaussian noise \epsilon을 섞어 noisy action A_t^\tau를 생성합니다. Gaussian noise는 정규분포에서 sampling하고, \tau는 flow matching timestep인데, beta 기반 분포에서 sampling하여 작은 값을 가지도록 설계하여 noise에 가까운 값이 더 자주 나오도록 설계했습니다.
\epsilon \sim \mathcal{N}(0, I) , A_t^\tau = \tau A_t + (1-\tau)\epsilon
loss는 vector field에서 모델의 예측 v_{\theta}(A_t^\tau, o_t)과 정답 u(A_t^\tau \mid A_t)의 차이를 최소화합니다. GT action 자체를 출력하는 것이 아니라 복원하기 위한 방향을 예측합니다. 이렇게 설계한 이유는 추론에서 단계별로 적분하기 때문인 것 같습니다.
\mathcal{L}^{\tau}(\theta) = \mathbb{E}_{p(A_t \mid o_t),\, q(A_t^{\tau} \mid A_t)} || v_{\theta}(A_t^{\tau}, o_t) - u(A_t^{\tau} \mid A_t) ||^2

추론에서는 random noise에서 시작하여 \tau = 0 에서 \tau = 1 까지 \delta = 0.1 간격으로 10번 적분하여 action을 생성합니다.
A_t^{\tau+\delta} = A_t^\tau+\delta v_\theta(A_t^\tau,o_t)

Data Collection and Training Recipe

저자는 robot foundation model을 위해 architecture 뿐 아니라 dataset과 training 방법도 중요하다고 합니다. Introduction에서 언급한 것처럼 pre-training / post-training을 진행하며, 어떤 데이터를 어떻게 학습하는지 설명합니다.

A. Pre-training and post-training

pre-training mixture는 timestep (o_t, A_t) 로 볼 수 있어서 timestep 단위로 측정한다고 합니다. pre-training mixture는 open-source dataset(9.1%)과 자체 수집한 dexterous manipulation dataset(나머지)으로 구성됩니다. open-source dataset은 넓은 범위의 물체와 환경을 포함하고, 자체 수집 데이터는 68개의 다양하고 복잡한 task로 구성됩니다. 그리고 task를 수행하면서 훨씬 다양한 행동 양상으로 데이터를 수집했다고 합니다. dataset 양은 불균형하기 때문에 보정을 위해 sample 수가 많은 task-robot 조합의 학습 비중을 줄이는 방법을 사용합니다. post-training에서는 더 작은 task-specific dataset으로 모델을 fine-tuning 합니다. task마다 필요한 dataset이 매우 다르기 때문에 각 task에 필요한 시간만큼 모았다고 합니다.

B. Language and high-level policies

복잡한 task는 바로 주지 않고 high-level VLM으로 decomposition을 진행하여 subtask로 나눠서 진행했다고 합니다. 이걸 입력으로 받아서 subtask의 action을 출력하도록 했다고 합니다.

C. Robot system details

저자는 pi-0가 하나의 robot에만 적용되는 모델이 아니라 7개의 다른 robot configuration 데이터를 함께 학습하는 cross-embodiment model이라고 합니다. 로봇마다 관절 개수나 action dimension 형태, camera 수가 달라도 공통 representation으로 맞춰 cross-embodiment training을 가능하게 했습니다.

Experimental Evaluation

저자는 pi-0를 4가지 측면으로 평가했습니다.

1. pre-training 후 여러 task를 얼마나 잘 수행하는가
2. language command를 얼마나 잘 이해하는가
3. 이전의 dexterous manipulation task에 비해 얼마나 잘 수행하는가
4. complex, multi-stage task에서 얼마나 잘 수행하는가

A. Evaluating the base model

post-training 없이 base model 만으로 여러 task를 zero-shot으로 얼마나 잘 수행하는가를 평가하기 위해 실험합니다.

평가 task는 shirt folding, bussing easy, bussing hard, grocery bagging, toast out of toaster이며, 모두 language command로 하나의 base model에서 진행했다고 합니다.

pi-0 full이 모든 task에서 가장 높은 성능을 보이며, shirt folding과 bussing easy에서는 거의 완벽에 가까운 성능을 보였습니다. 특히 다른 baseline 모델과 step 수를 맞춘 pi-0 parity도 모든 task에서 가장 높은 성능을 보이는데, 단순히 학습이 길어서 성능이 잘 나온 것이 아니라는 증거가 된다고 합니다. 저자는 OpenVLA가 action chunk와 high-frequency control을 지원하지 않는 autoregressive action discretization 구조이기 때문에 dexterous task에서 어려움을 겪는다고 언급합니다. Octo는 action chunk를 지원하지만 model capacity가 작아서 이런 결과가 나왔다고 합니다.

B. Following language commands

VLM pre-training이 language instruction following 능력을 얼마나 향상시키는지 확인합니다.

base model에 pre-training을 적용하고 table bussing, table setting, grocery bagging의 language-conditioned data로 fine-tuning하고, VLM backbone이 없는 pi-0-small과 비교합니다. 약 2초 길이의 language-labeled segment 여러 개가 이어져 하나의 task를 구성합니다. 평가는 전체 task command만 받는 flat condition, 사람이 중간에 step command를 제공하는 human condition, high-level VLM policy가 command를 생성하는 HL condition을 비교합니다.

pi-0는 pi-0-small보다 높은 language following accuracy를 보입니다. 또한 pi-0는 human condition에서 가장 큰 성능 향상을 보이고, HL에서도 flat보다 성능이 향상됩니다. 반면 pi-0-small은 high-level guidance를 추가해도 성능 향상이 제한적입니다. 저자들은 이를 통해 pre-training이 language instruction을 실제 robot action으로 전환하는 능력에 기여한다고 결론내립니다.

C. Learning new dexterous tasks

pre-training data와 크게 다른 새로운 dexterous task를 얼마나 잘 학습하는지 평가합니다. 저자들은 각 task에 서로 다른 양의 fine-tuning data를 사용했으며, pre-training task와의 유사도에 따라 task를 easy와 hard로 구분했습니다.

진행한 task에서 stack bowls와 towel folding은 pre-training task와 유사한 easy task로 구분합니다. microwave task는 microwave라는 새로운 요소를 포함하고, Paper towel replacement와 Franka items in drawer는 pre-training에 유사한 물체나 robot-task 조합이 없어서 hard task로 구분했다고 합니다.

저자는 OpenVLA와 Octo는 OXE pre-trained에서 fine-tuning하고 ACT와 Diffusion Policy는 task-specific fine-tuning data만으로 처음부터 학습했습니다. pi-0는 pre-trained base model에서 fine-tuning한 경우와 task data만으로 scratch training한 경우를 모두 평가합니다. 이렇게 한 이유는 pi-0의 architecture와 pre-training의 효과를 구분해서 확인하려고 했다고 합니다.

저자는 task 별로 다른 양의 fine-tuning data를 사용하여 실험 결과를 나타냈습니다. 전반적으로 pi-0가 높은 성능을 보이며, 1시간의 microwave fine-tuning data를 사용한 pi-0가 baseline보다 상당히 높게 나타난 것을 확인할 수 있습니다. pre-training 효과 측면에서는 pre-training data와 더 유사한 task에서 더 크게 나타났다고 합니다. pre-trained는 scratch 보다 대부분 더 높은 성능을 보이고, 일부 task에서는 2배까지 높았다고 합니다.
이런 결과를 통해서 저자는 pre-training이 새로운 dexterous task를 더 적은 data로 학습해도 된다는 결과를 도출했다고 합니다. 그리고 새로운 task가 pre-training에 있는 task와 얼마나 유사한지도 영향을 받는다는 것을 알아냈다고 합니다.

D. Mastering complex multi-stage tasks

fine-tuning과 language를 결합해서 긴 시간 동안 여러 행동을 순서대로 수행해야 하는 complex multi-stage task에 적응할 수 있는지 평가합니다.

일부 task는 pre-training data에 이미 포함되어 있지만 높은 완성도를 위해 fine-tuning이 필요하고, 일부 task는 pre-training에 전혀 없는 새로운 task입니다. 빨래 관련 task는 pre-training에 포함되어 있지만, table bussing, box building, to-go box, egg packing은 pre-training에 없는 task 입니다. 특히 새로운 task는 변형 가능한 옷이나 cardboard, 깨지기 쉬운 계란, 조밀한 clutter, 양팔 협응, 여러 단계의 행동 순서 결정과 같은 어려움이 있다고 합니다.

결과를 보면 각 task에서 10 trial 평균 score를 나타냈습니다. 다른 방법들에서는 이 task를 해결하지 못했기 때문에 pre-training 후 fine-tuning한 full, pre-training만 수행한 zero-shot, 그리고 fine-tuning data만으로 학습한 scratch만 비교했다고 합니다. 결론적으로 full은 모든 task에서 0.5 이상의 점수를 보였고, 전반적으로 zero-shot과 scratch보다 좋은 성능을 보였습니다. 특히 더 복잡한 task에서 pre-training의 성능 향상 폭이 큰 것을 확인하여 다양한 robot data로 pre-training 한 뒤 특정 task에 high-quality data로 fine-tuning 하는 학습 방법이 complex, multi-stage dexterous manipulation을 수행하는 것에 중요한 역할을 한다고 주장합니다.

Discussion, Limitation, and Future Work

저자들은 pi-0에 대해서 다양한 robot data로 pre-training한 뒤 complex downstream task의 fine-tuning을 수행하는 robot foundation model framework라고 다시 정리했습니다. VLM pre-training은 language와 visual semantic knowledge를 제공하고, flow matching은 high-frequency continuous action chunk를 생성합니다. 실험은 pre-training이 broad capability를 만들고, post-training이 복잡한 task의 fluent한 수행 전략을 학습하게 한다는 점을 보인다고 해석합니다.

저자는 pre-training dataset을 어떤 구성으로 만들고, 어떤 data에 어떤 weight를 부여해야 하는지 이해가 부족하다고 합니다. 또한 모든 evaluation task가 안정적으로 성공하지는 않으며, 특정 task에서 near-perfect performance를 얻기 위해 필요한 data의 양과 종류를 예측하기 어렵다는 것도 지적했습니다. 서로 다른 task와 robot 데이터를 함께 학습할 때 발생하는 positive transfer의 범위도 명확하지 않다는 것을 limitation으로 말합니다.

어떤 pre-training data가 실제로 도움이 되는지, data mixture를 어떻게 weighting해야 하는지, cross-embodiment pre-training의 효과가 manipulation을 넘어 autonomous driving, navigation, legged locomotion처럼 더 다른 domain까지 확장되는지 future work로 남겨두면서 마무리 했습니다.

---

마치며

연구실 생활을 하면서 느끼기도 하지만, 이 논문을 보면서 로봇 분야와 다른 분야의 관계가 다 엮여있다는 느낌을 많이 받았습니다. 로봇이 기술의 집약체인 만큼 vision이나 language 분야의 발전으로 로봇 분야도 새로운 길이 열리는걸 잘 보여준 케이스라고 생각합니다. 다양한 실험으로 받아들이면서 읽었는데 비판적으로 읽지 못한 느낌도 들어서 그런 점은 좀 아쉬운 것 같습니다. 이번 리뷰를 통해 다음 pi 시리즈를 읽거나 pi 시리즈를 언급하면서 진행되는 연구들과 대화들에 많은 도움이 될 것 같습니다. 감사합니다.