[arXiv 2024] π0: A Vision-Language-Action Flow Model for General Robot Control

이번 리뷰 논문은 VLA 논문으로 지능형 로봇을 선두하고 있는 기관 중 Physical Intelligence라는 곳에서 공개한 기법입니다. 최근 VLA의 흐름은 사전 학습된 VLM 모델에 Diffusion policy를 붙이는 방법들이 주축을 이루고 있는 것 같습니다. 이러한 흐름에 큰 기여를 한 기법으로 확인됩니다.
+ Diffusion Policy을 붙이는 이유는 아마 로봇의 액션이 연속적인 공간에서 표현되기 때문인 것 같습니다. 이전 세미나에서 언급한 바와 같이 이야기했지만 Diffusion은 GAN과 같은 생성형 모델과 다르게 generative model에 속합니다. 이러한 특징이 학습되는 데이터 공간들을 연속적으로 연결하도록 해주는 큰 요소라고 볼 수 있습니다. 이러한 특징 때문에 diffusion을 이용하는 것 같습니다. (diffusion 하면 실시간이 떨어지기 때문에 로봇에 적합한가? 라는 궁금증이 생길 수 있습니다. 로봇 액션의 차원 (e.g. 연구실 보유 중인 매니퓰레이터 UR5e는 6 joints + gripper로 7차원의 벡터만 요구됨)과 영상 차원은 상대적으로 차이가 많기 때문에 차원의 저주에 대한 영향 덜합니다. 즉, 상대적으로 적은 iterations과 작은 모델로도 필요한 값을 생성할 수 있어 추론 속도 문제는 완화됩니다.)

Intro

2016년 알파고와 이세돌의 대결을 기억하시는 분들이 있을 겁니다. 이때를 기점으로 인공지능에 대한 붐이 일어나 고도로 전문화된 시스템부터 인간의 사고로는 접근하기 힘든 복잡한 문제를 해결하는 시스템, 예를 들어 단백질의 형태를 예측한 AlphaFold는 과학계의 공로를 크게 인정 받아 화학자가 아닌 컴퓨터 공학자가 개발한 기술임에도 불구하고 2024년도 노벨 화학상을 받게 되었습니다. 이외에도 저희 연구실의 연구원분들은 ChatGPT를 잘 활용하시고 계신 것을 알고 계십니다. 저도 그렇지만 논문을 읽다가 어려운 부분이 생기거나, 글을 작성하는 데에 있어 창의적인 발상에 도움을 받거나 글의 어투를 수정하기 위해서 많이 사용합니다. 즉, 이미 인공지능은 저희의 일상 속에 침투하여 인간의 일부 능력을 뛰어 넘어 도움을 주고 있습니다.

허나, 인공지능이 아직 사람을 능가하지 못하고 있는 영역이 있습니다. 바로 “다양성” 입니다. 특히 다양한 물리적 환경에서 적응적으로 다양한 과제를 해결하는 능력은 인간을 뛰어넘지 못하고 있습니다. 최근 연구 중에서 이러한 한계를 극복하기 위한 발전을 보이는 대표적인 연구가 VLM이라고 볼 수 있습니다. VLM은 인터넷 스케일의 대규모 데이터 셋으로 사전 학습된 시스템으로 언어 지시와 영상 기반의 장면에 대한 정보를 추상적으로 유추하여 지시에 따른 뛰어난 문제 해결 능력을 보여주었습니다. 하지만 VLM은 사람과 같은 물리적 세계에 위치한 정보들을 직접 얻는 것이 아니라 물리적 상호작용에 대한 이해를 전적으로 추상적인 추론 능력에 의존합니다. 그렇기에 물리적인 정보에 대한 추론 능력은 현저하게 떨어진다는 문제가 있습니다.

저자는 VLM에게 물리적으로 위치한 정보들을 직접 학습을 시켜 사람과 동일한 물리적 공간에서의 문제 해결 능력을 갖추는 방법을 제안합니다. 사실 물리적인 공간에 대한 이해를 시키기는 건… 문제 복잡도가 정말 미치도록 높다는 문제가 있죠. 다행히도 최근 인공지능에는 차용 가능한 기술들이 등장하기 시작했습니다. 대표적으로 VLM, LLM, Foundation Model.들은 인터넷 스케일의 대용량 데이터 셋으로 사전 학습하고 특정 태스크에서 평가 (zero-shot 혹은 fine-tunning)을 수행하였을 때, 평가 데이터 셋에 특수하게 학습된 모델들을 뛰어 넘는 결과를 보여주었습니다. 저자는 이러한 방법을 착안하여 논리적인 추론 능력을 유지하기 위해서 사전 훈련된 VLM을 이용하면서 인터넷 스케일에 준하는 로봇 에이전트 데이터 셋을 이용하여 물리적인 정보에 대한 직접적인 문제 해결 능력을 부여할 것을 제안합니다.

저자는 이를 구현하기 위해서는 3가지 요소들(1. very large scale data. 2. right model architectures, 3.right training recipe)을 해소해야 한다고 주장합니다. 저자는 이러한 요소들을 해소한 π0를 제안합니다. 앞서 언급하 바와 같이 인터넷 스케일의 데이터 셋 사전 학습된 VLM을 이용합니다. 이는 물리적인 환경에서의 일반화된 지식, 의미론적인 사고, 문제 해결 해결 능력을 계승 받는 것을 목적으로 합니다. 추가로 Cross-embodiment training을 채용하여 다양한 로봇 타입, 다양한 태스크, 다양한 구성에 일반화가 가능하도록 대용량 데이터 셋 Open X-Embodiment (OXE) 데이터 셋과 long-term task를 포함한 저자가 구축한 데이터 셋이 결합된 대용량 데이터 셋으로부터 물리적인 문제 해결 능력을 부여 받습니다.

더 나아가, 정교하고 복잡한 물리적은 태스크를 수행 가능하도록 하기 위해서 VLM에 추가적으로 flow matching (variant of diffusion)을 따르는 action chunking architecture를 추가합니다. (flow matching은 복잡하고 연속적인 액션 분포를 표현하기 위해서 사용합니다.) 해당 구조는 빨래 접기와 같은 태스크 (fig 2)에서의 액션 정보에 대한 출력을 50hz까지 가능하게 만듭니다. (fig 1의 action expert가 action chunking architecture에 해당합니다.)

마지막으로 저자는 복잡한 태스크를 유연하고 강인하게 수행하기 위해서 사전 학습과 사후 학습에 대한 데이터 셋을 구분하여 진행해야 한다고 합니다. 사전 학습 데이터 셋에서는 거친 데이터 셋을 통해 노이즈에 강인하면서 일반화된 능력을 확보하고 사후 학습 데이터 셋에서는 복잡하고 정교한 태스크를 정확하게 수행하기 위해서 잘 정제된 데이터 셋으로부터 학습을 진행해야 한다고 합니다.

이렇게 구성된 π0는 최신 기법 대비 zero-shot, fine-tunning, long-term 에서 모두 SOTA 성능을 보여줍니다.

Method

+ 해당 논문에 대해서 아쉬운 점이 있습니다… Method의 알맹이들이 이전 기법들에서 가져온 것이기에 자세한 설명은 생략하고 있습니다. 각 방법론들이 굵직한 기법이기에 간단하게 설명하기에는 주객이 전도하기 때문에 생략한 것 같습니다. 이에 대해서는 제가 얼른 관련 기법들에 대해서 리뷰를 작성하도록 하겠습니다.

Overview. 전반적인 학습 프레임워크는 fig 3에서 확인이 가능합니다. 먼저, 정교한 작업이 포함된 저자가 제안한 데이터 셋 (7 different robot configurations for 68 different tasks)과 OXE datasets (22 robots) 간의 가중 결합으로 사전 학습 데이터 셋을 볼 수 있습니다. 그 다음, 사전 학습 단계에서는 task names과 segment anotations*으로 구성된 다양한 언어 라벨을 이용합니다. 사전 학습 단계에서는 폭 넓은 다양성과 일반화된 능력을 보여주는 base model을 학습하는 것이지 특정 작업에서 높은 성능을 달성하기 위해서 특화되는 것이 목적이 아닙니다. 복잡한 작업을 수행하기 위해서 post-training phase를 추가적을 수행합니다. (아쉽게도 사후 훈련에 대해서 새로운 것을 제안하지는 않습니다.)
저자가 제안한 모델은 PaliGemma라는 VLM 모델에 flow mathing을 추가된 형태 입니다. 여기서 PaliGemma는 비교적 작은 크기이기에 로봇에 적합하다고 저자가 판단하여 채용된 것이지 다른 VLM을 이용해도 무방하다고 주장합니다.
+ segment anotations*: fine-grained labels로 2초 길이로 되어 있다고 함 (실제 데이터 셋 확인 필요)

Flow matching. π0의 모델 구조는 기본적으로 Transfusion을 따릅니다. (Transfusion은 VLM과 영상 생성 모델을 결합한 모델로 VLM + Diffusion model이라고 보시면 됩니다.) 먼저, 이미지 인코더로 로봇이 관측한 영상들을 언어 토큰과 동일한 임베딩 공간에 사영 시킵니다. 그 다음, 로봇의 액션에 대한 정보들을 diffusion policy을 기반으로 하여 가우시안 분포를 따르는 액션 초기값을 입력 시켜 최종적인 출력값을 얻습니다. (Transfusion의 diffusion model 부분이 diffusion policy로 변경된 것으로 이해하시면 됩니다) 이는 mixture of experts로 VLM을 확장한 기반 연구들을 기반으로 합니다. first expert는 영상과 텍스트로 구성되며 second expert는 로보틱하게 특화된 정보들이 입력됩니다. 저자는 이를 action expert라고 지칭합니다.

구체적으로 저자는 data distribution p(A_t|o_t) 를 따르는 조건부 분포를 생성하기를 바랍니다. 여기서 A_t = [a_t, a_{t+1}, ... , a_{t+H-1}] 이며, H는 action chunk의 크기로 미래의 행동으로 구성됩니다. (저자는 H=50으로 세팅) t 시점에서 로봇의 관측 상태 o_t = [I_1^t, …, I_n^t, l_t, q_t] 는 t 시점에서의 I_1^t, …, I_n^t 로 로봇이 수집한 n개의 영상들과 “빨래를 접어라”와 같은 언어 지시 l_t , 로봇의 액션 정보~joint angles q_t 로 구성됩니다. 영상과 로봇의 액션 정보들은 각각의 encoder를 통해 인코딩된 다음 linear projection layer를 통해 텍스트와 동일한 임베딩 공간에 사영 됩니다.
+ 여기서 영상들은 대체로 로봇당 2~3개 정도로 구성됩니다. 로봇을 관측하는 카메라와 로봇 팔에 탑재된 카메라로 구성됩니다.

학습 중에는 action tokens을 아래와 같은 conditional flow matching loss를 통해 지도학습을 합니다.

여기서 아래 첨자의 t는 robot timesteps, 윗 첨자는 flow matching timesteps τ ∈ [0, 1]을 따릅니다. 앞서 언급한 바와 같이 flow matching은 diffusion model과 VLM을 결합한 영상 생성 모델 Transfusion을 기반으로 합니다. 이에 따라, 간단한 linear-Gaussian probability path를 기반으로 하여 q(A_{t}^{\tau} | A_t) = N(\tau A_t, (1 - \tau)I) 를 따르도록 합니다. (여기서 I는 단위 행렬입니다.) 즉, 정규 분포를 따르는 random noise \epsilon \sim N(0, I) 를 샘플링하여 “noisy actions” A_{t}^{\tau} = \tau A_t + (1 - \tau)\epsilon 을 점진적으로 생성하여 diffusion model v_\theta(A_{\tau t}, o_t) 를 학습하는 방식이라고 보면 됩니다. 상단 수식에서 u(A^{\tau}_t | A_t) = \epsilon − A_t으로 최종적으로 모델은 \epsilon 만 예측하면 됩니다. 그럼 최종적은 예측은 아래 수식으로 얻을 수 있게 됩니다.

여기서 δ는 step size로 저자는 10사용 했기에 0.1로 구성됩니다. 추론 단계에서는 영상과 텍스트에 대한 attention key와 values는 캐싱되어져 iteration을 수행하면 되기 때문에 추론 속도 개선에 도움을 줄 수 있습니다.
+ 더 심플하게 설명하면 액션 추론만 diffusion으로 예측하기 때문에 위 수식으로 반복 연산을 수행해야하지만 이 때 사용되는 영상과 텍스트는 같은 값을 사용해도 무방하다는 뜻입니다.
+ 3 billion parameter를 가진 PaliGemma를 기반으로 300M parameters를 가진 action expert를 추가
++ Non-VLM baseline model: VLM의 영향을 실험하기 위해서 VLM이 없이 구성된 π0-small로 추가 실험함. 해당 모델은 470M parameters를 가짐

DATA COLLECTION AND TRAINING RECIPE

fig 4와 같이 로봇 구성에 따른 데이터 분포가 불균형이 심하기 때문에 저자는 n^{0.43} 의 가중치를 부여합니다. 여기서 n은 샘플의 수로 샘플 수가 많은 경우에는 낮은 가중치를 부여 받게 됩니다. 이렇게 구성된 데이터 셋이 fig 4의 오른쪽에 해당합니다.

또한 configuration vector q_t와 action vector a_t는 항상 데이터 셋 구성에서 가장 큰 값을 따릅니다. 즉 18개 (two 6DoF arms, 2 grippers, mobile base, vertically actuated torso)를 가지며 나머지는 zero-pad로 채워집니다. 영상 또한 이를 따르며 3 장의 영상을 입력으로 부여되며, 3장 이하인 경우에는 부족한 영상에 대한 token을 zero-pad로 채웁니다.

사후 훈련 단계에는 간단한 태스크인 경우에는 5시간, 복잡한 경우에는 100시간 이상으로 구성된 데이터를 구축하여 실험을 진행합니다.

Language and high-level policies. 복잡한 태스크인 경우에는 의미론적인 추론 능력과 high-level strategy를 요구하게 됩니다. 저자는 이를 해소하기 위해서 high-level VLM을 추가적으로 활용해 high-level policy를 수행할 것을 제안합니다. 이에 대한 실험 또한 포함 됩니다.

Robot system details. 정보의 나열이기 캡쳐본으로 대처하도록 하겠습니다. 로봇 구성은 fig 5에서 확인 가능합니다.

Experiment

Evaluating the base model

해당 실험에서는 사전 학습이 얼마나 큰 효과를 보여주는지를 보기 위한 실험으로 사후 훈련 없이 zero-shot으로만 평가를 진행합니다. 진행된 결과는 fig 7에서 보이며, 수행한 태스크는 fig 6에서 확인 가능합니다. 태스크에 대한 구체적인 설명은 아래 캡쳐를 확인하시길 바랍니다.

비교를 수행하기 위한 모델은 7B의 파라미터로 OXE datasets으로 사전 학습된 OpenVLA와 93M의 OCto로 구성됩니다. Octo는 VLA는 아니지만 diffusion process 기반인 방법론 입니다. Octo인 경우, 저자가 혼합 구성한 데이터 셋으로 재학습을 진행했다고 합니다. 추가로 fig 7에서 “parity”를 확인 할 수 있습니다. 이는 160k steps으로 학습된 OpenVLA와 320k로 학습된 Octo로 비교하기 위해 700k로 학습된 모델을 160k로 학습한 모델입니다. OpenVLA (UR5e only)는 다양성을 배제하고 UR5e에서만으로도 능력을 확인하기 위해서 학습된 모델입니다.

놀랍게도 fig 7에서 보이는 바와 같이 π0 뿐만이 아니라 parity, small 조차도 우세한 결과를 보여줍니다. 이는 VLM에 대한 사전 지식이 도움이 됨을 증명할 뿐만이 아니라, 제안한 action expert의 우수함을 보입니다.

Following language commands. 해당 실험에서는 해당 모델이 얼마나 언어 지시를 잘 이해하는 지에 대한 실험으로 구성됩니다.

해당 실험에서 “flat”을 태스크 자체를 라벨로 부여한 결과이고, Human은 사람이 직접 구체적인 언어 지시를 부여한 결과, HL은 high-level VLM을 이용하여 구체적인 언어 지시를 생성하여 평가한 결과 입니다. 결과적으로 fig 9의 왼쪽 같이 VLM의 사전 지식이 언어를 이해하는 능력에 큰 도움을 준다는 것을 보여주며 오른쪽에서는 이를 더욱 극단적으로 보여줍니다. 사람이 구체적으로 지시에 대한 도움을 주는데에도 불구하고 π0-small-human에서 더 낮은 결과를 보여줌에 따라 VLM을 이용하는 것이 언어 지시에 대한 이해 능력에 큰 도움을 주는 것을 증명합니다.

Learning new dexterous tasks. 해당 실험에서는 사후 훈련을 통해 복잡한 작업을 얼마나 잘 대응할 수 있는지를 보는 실험으로 태새크의 난이도를 “tier”로 구분하였습니다. 티어는 학습한 데이터 셋에 유사한 작업이 있음과 만약에 있다면 얼마나 다른지를 감안하여 부여하였습니다. 해당 내용은 아래 캡쳐본을 참고 하시길 바랍니다.

결과는 fig 11에서 확인이 가능합니다. 거의 대부분의 실험에서 우세한 성능을 보여주었으며, 대체로 scratch 대비해서 더 좋은 성능을 보여주고 있습니다. 이는 VLM의 사전 지식을 유지하는 것이 더 좋다는 것을 방증합니다.

Mastering complex multi-stage tasks. 해당 실험에서는 여러 단계로 구성된 태스크를 얼마나 잘 수행하는 가를 평가한 실험입니다. 태스크에 대해서는 캡쳐본 첨부하겠습니다.

실험 결과, fig 13과 같이 사후 훈련 데이터 셋으로 fine-tunning을 수행한 결과가 가장 좋은 결과를 보여주면서 저자가 주장하는 사전 학습 데이터가 가지는 장점을 증명합니다.

---

해당 논문을 좀 더 자세하게 이해하려면 저자가 사용한 기법들에 대해서 공부를 더 해야 할 것 같습니다.