오늘은 Physical Intelligence에서 공개한 Black, Kevin, et al. “π0: A Vision-Language-Action Flow Model for General Robot Control” 논문을 리뷰해보겠습니다. 최근 로봇 분야에서 VLM(Vision-Language Model)에 Action을 결합한 VLA(Vision-Language-Action) 모델이 굉장히 많이 연구되고 있는데요. π0는 여기에 Flow Matching 을 적용하여 연속적인 로봇 Action을 생성하는 모델입니다.

이번 글에서는 제가 논문을 읽으면서 중요하다고 생각한 Introduction, Overview, π0 Model 을 중심으로 이야기해보겠습니다. Related Work와 세부 실험 결과는 생략하고, π0가 왜 나왔고 기존 VLA와 무엇이 다른지, 그리고 핵심인 Flow Matching 수식을 최대한 천천히 살펴보겠습니다. 수식에서 제가 공부하면서 헷갈렸던 부분도 마지막에 따로 정리해두었습니다😎

Overview of pi0

[그림 1] Overview of π0

Introduction

논문은 Robert A. Heinlein의 “Specialization is for insects.” 라는 문장으로 시작합니다. 말 그대로 인간 지능의 중요한 특징은 하나의 일만 잘하는 것이 아니라 굉장히 다양한 일을 수행할 수 있다는 뜻입니다.

논문에서는 인간과 현재 AI의 차이를 설명하면서 Versatility(범용성) 을 강조합니다. 사람이 가진 범용성은 단순히 여러 Task를 수행한다는 의미가 아니라,

  • 다양한 물리 환경에서 여러 Task 수행

  • 환경의 제약에 대응

  • 자연어 명령 이해

  • 예상하지 못한 상황이나 Perturbation에 대응

까지 포함하는 개념입니다.

최근 LLM과 VLM은 굉장히 큰 데이터로 Pre-training한 뒤, 잘 정제된 데이터로 Fine-tuning 또는 Alignment를 수행하면서 이러한 범용성에 가까워지고 있습니다. 하지만 VLM에는 중요한 문제가 있습니다.

VLM은 실제 물리 세계에 존재하지 않습니다.

VLM은 이미지와 텍스트를 통해 “셔츠를 접는 방법”은 알고 있을 수 있지만, 실제 로봇이 셔츠를 잡을 때 어느 지점을 잡아야 하는지, 천이 미끄러졌을 때 어떻게 복구해야 하는지, 얼마만큼 관절을 움직여야 하는지는 직접 경험하지 못했습니다. 즉, 물리적 상호작용에 대한 이해가 대부분 추상적인 설명에 기반합니다.

그래서 π0에서는 인터넷 규모의 지식을 가진 Pre-trained VLM 에 실제 Embodied Robot Data 를 결합합니다.

\[\text{Internet-scale VLM Knowledge} + \text{Robot Interaction Data}\]

이러한 방식은 로봇 학습에서 크게 세 가지 문제를 해결할 가능성이 있습니다.

  1. Data Scarcity : 하나의 Task 데이터만 쓰지 않고 다른 Task, 다른 Robot의 데이터까지 활용할 수 있음

  2. Generalization : 더 다양한 환경과 상태를 학습하기 때문에 새로운 상황에 대응할 가능성이 높음

  3. Robustness : 다양한 데이터 안에 Correction, Recovery behavior가 존재하기 때문에 실수 이후 복구도 배울 수 있음

즉 기존 방식처럼 Task A를 위한 Policy A, Task B를 위한 Policy B를 각각 만드는 것이 아니라, 다양한 데이터를 먼저 학습한 Robot Foundation Model 을 만들고 원하는 Task에 맞게 사용하는 것입니다.

\[\text{Diverse Robot Data} \rightarrow \text{General Robot Foundation Model} \rightarrow \text{Downstream Task}\]

Robot Foundation Model의 세 가지 Bottleneck

그렇다면 그냥 데이터를 많이 넣고 큰 모델을 만들면 끝일까요?? 논문에서는 Robot Foundation Model을 만들기 위해 해결해야 할 세 가지 Bottleneck을 이야기합니다.

  1. Scale

    Foundation Model의 장점은 작은 규모에서는 잘 나타나지 않을 수 있기 때문에 충분히 큰 데이터와 모델 규모가 필요합니다.

  2. Architecture

    이미지와 언어의 Semantic Information을 잘 이해하면서 동시에 굉장히 정교한 Continuous Robot Action까지 표현해야 합니다.

  3. Training Recipe

    논문에서는 이 부분이 어쩌면 가장 중요하다고 이야기합니다. LLM처럼 어떤 데이터를 Pre-training에 쓰고 어떤 데이터를 Post-training에 사용할 것인지가 중요합니다.

π0는 이 세 가지를 해결하기 위해 다음의 요소들을 사용합니다.

  • Pre-trained VLM
  • Cross-embodiment Training
  • Action Chunking
  • Flow Matching
  • Action Expert
  • Pre-training / Post-training Recipe

하나씩 살펴보겠습니다.

VLM에서 VLA로

π0는 처음부터 모든 것을 학습하지 않고 이미 인터넷 규모의 이미지와 언어를 학습한 VLM을 사용합니다. 이를 통해 VLM이 가지고 있는 일반 지식, Semantic Reasoning, Problem Solving 능력을 가져옵니다.

그리고 여기에 Robot Action을 추가로 학습하면서,

\[\text{Vision} + \text{Language} \rightarrow \text{Vision} + \text{Language} + \text{Action}\]

즉 VLM을 VLA로 확장 합니다.

또 하나 중요한 점은 Cross-embodiment Training 입니다. 하나의 로봇 데이터만 이용하는 것이 아니라 Single-arm, Dual-arm, Mobile Manipulator처럼 서로 다른 Configuration Space와 Action Representation을 가지는 로봇들의 데이터를 하나의 모델에서 학습합니다.

논문에서는 7개의 서로 다른 Robot Configuration과 68개의 Task에서 자체 데이터를 수집했고, 여기에 여러 로봇으로 구성된 Open X-Embodiment(OXE) 데이터도 함께 사용합니다.

pi0 laundry folding

[그림 2] π0를 이용한 Laundry Folding

그림처럼 π0는 단순히 물건 하나를 집는 Task만 수행하는 것이 아니라 Dryer에서 빨래를 가져오고, Hamper에 넣고, 테이블로 이동한 뒤 여러 옷을 접는 Long-horizon Task까지 목표로 합니다. 논문에서는 전체적으로 10,000시간이 넘는 Robot Data를 이용합니다.

Pre-training & Post-training

π0를 이해할 때 모델 구조만큼 중요한 것이 Pre-training과 Post-training을 분리한다는 점 입니다.

이 아이디어는 LLM의 학습 방법과 상당히 비슷합니다.

구분 Pre-training Post-training
데이터 크고 다양함 상대적으로 작고 정제됨
목표 Broad Capability, Generalization Specific Task의 높은 숙련도
행동 다양한 상황과 Recovery 포함 효율적이고 일관된 행동
결과 Base π0 Specialized π0

논문의 설명이 재미있는데요. 고품질 데이터만 학습하면 오히려 Recovery를 배우기 어렵습니다. 고품질 Demonstration에는 애초에 실수가 잘 등장하지 않기 때문입니다.

반대로 다양하지만 상대적으로 품질이 낮은 데이터만 사용하면 여러 상황은 경험할 수 있지만 깔끔하고 효율적으로 Task를 수행하는 방법을 충분히 배우기 어렵습니다.

그래서,

\[\text{Diverse Pre-training} + \text{High-quality Post-training}\]

을 함께 사용합니다.

쉽게 이야기하면 Pre-training에서 이것저것 경험해보고, Post-training에서 잘하는 방법을 다듬는다!! 정도로 생각하면 될 것 같습니다.

Overview

Architecture of pi0

[그림 3] Overall Framework of π0

이제 본격적으로 전체 구조를 보겠습니다. [그림 3]만 잘 이해해도 π0의 전체적인 흐름을 거의 다 이해할 수 있습니다.

먼저 Pre-training Mixture에는 자체적으로 수집한 Dexterous Manipulation Dataset과 OXE가 들어갑니다. 자체 데이터는 7개의 Robot Configuration, 68개의 Task로 구성되어 있고, OXE에는 여러 종류의 로봇 데이터가 포함됩니다.

또한 Language Label을 단순히 전체 Task 이름으로만 주지 않습니다.

  • Task Name

  • Segment Annotation

두 종류를 사용합니다. Segment Annotation은 일반적으로 약 2초 길이의 Sub-trajectory에 붙는 Fine-grained Label입니다. 즉 “빨래 접기”라는 큰 명령만 주는 것이 아니라 긴 Trajectory 안의 세부 행동에도 Language Grounding을 제공하는 것입니다.

Pre-training을 끝낸 Base Model은 언어 명령을 따르고 다양한 Task를 수행할 수 있지만 논문에서는 이를 rudimentary proficiency, 즉 기초적인 숙련 수준이라고 표현합니다. 여기서 특정 어려운 Task를 더 잘 수행하고 싶다면 High-quality Data를 이용해 Post-training을 합니다.

논문에서는 Post-training도 크게 두 가지 형태를 이야기합니다.

  • Efficient Post-training : 적은 양에서 중간 정도의 데이터로 새로운 Task에 빠르게 적응

  • High-quality Post-training : Laundry Folding, Mobile Manipulation과 같은 어려운 Task에 더 많은 고품질 데이터 사용

전체 흐름을 간단하게 표현하면 아래와 같습니다.

\[\text{Diverse Pre-training} \rightarrow \text{Base }\pi_0 \rightarrow \text{High-quality Post-training} \rightarrow \text{Dexterous Specialist}\]

π0 Model

이제 논문에서 가장 중요한 Section IV의 모델 구조를 살펴보겠습니다.

π0의 Backbone은 PaliGemma 입니다. PaliGemma는 대략 3B 크기의 VLM이고 [그림 3]에서는 SigLIP(400M) + Gemma(2.6B) 로 표현되어 있습니다.

π0는 이 VLM에 Robotics-specific Input과 Output을 추가합니다.

  • Vision

  • Language

  • Proprioceptive State

  • Robot Action

이미지는 Image Encoder를 통해 Language Token과 같은 Embedding Space로 Projection되고, Robot State 역시 Encoder와 Linear Projection을 통해 같은 공간에서 처리됩니다.

Observation과 Action Chunk

π0가 모델링하려는 것은 단순한 하나의 Action이 아니라 다음과 같은 조건부 분포입니다.

\[p(A_t \mid o_t)\]

여기서 Observation은,

\[o_t = [I_t^1, \cdots, I_t^n, \ell_t, q_t]\]

으로 정의됩니다.

  • \(I_t^i\) : i번째 RGB 이미지, 로봇마다 2~3개 사용

  • \(\ell_t\) : Language Token Sequence

  • \(q_t\) : Robot의 Joint Angle Vector

그리고 Action은 하나가 아니라 미래 여러 Action을 묶은 Action Chunk 입니다.

\[A_t = [a_t, a_{t+1}, \cdots, a_{t+H-1}]\]

논문에서는 \(H=50\) 을 사용합니다.

기존 Policy가 매 순간 \(a_t\) 하나를 생성한다고 생각하면, π0는 미래의 짧은 Action Sequence를 한 번에 모델링한다고 볼 수 있습니다.

여기서 한 가지 주의할 점!! 논문에서는 Dexterous Task를 최대 50Hz로 제어한다고 이야기합니다. 따라서 단순 계산으로 50개의 Action Chunk는 50Hz 기준 약 1초 분량에 해당할 수 있습니다. 다만 H=50이라고 해서 항상 50개 Action을 모두 실행한 뒤 다시 추론한다는 뜻은 아닙니다. Action Chunk의 생성 길이와 실제 Receding-horizon Execution 방식은 구분해서 보는 것이 좋습니다.

Action Expert

기존 VLM은 Image와 Text를 다루는 모델입니다. 따라서 Robot State와 Continuous Action을 그대로 잘 처리한다고 보기는 어렵겠죠?? π0에서는 Robotics-specific Token을 처리하기 위한 별도의 Weight Set을 추가하고 이를 Action Expert 라고 부릅니다.

논문에서는 이 구조를 두 개의 Expert를 가지는 Mixture-of-Experts(MoE)와 비슷하다고 설명합니다.

  • VLM Weight : Image / Text 처리

  • Action Expert Weight : Robot State / Action 처리

하지만 일반적인 MoE에서 Router가 Token마다 Expert를 선택하는 구조와 완전히 같다고 생각하면 안 됩니다. 여기서는 Robot-specific Token에 별도의 Weight Set을 사용하는 구조가 MoE와 유사하다는 의미에 가깝습니다.

또 Action Expert는 300M Parameter이고 Scratch부터 학습합니다. 결과적으로,

\[3.0B\;\text{PaliGemma} + 0.3B\;\text{Action Expert} = 3.3B\;\pi_0\]

정도의 모델이 됩니다.

추가로 VLM Pre-training의 효과를 확인하기 위해 VLM Initialization을 사용하지 않는 π0-small(470M) 도 Ablation용으로 학습합니다.

Action Token은 Discrete Token이 아니다!!

이 부분은 논문을 읽으면서 꼭 주의해야 합니다.

π0 논문에서 말하는 Action Token은 Language Token처럼 Action을 Discretization한 Token이라는 뜻이 아닙니다.

논문에서는 Token이라는 단어를 Sequence Dimension상의 Input/Output Slot이라는 의미로 넓게 사용합니다. 따라서,

  • Language Token : Discrete Variable
  • Image Patch : Continuous Variable
  • Robot Action Token : Continuous Variable

이 모두를 Token이라고 부릅니다.

이 점이 기존 Autoregressive VLA와 π0의 큰 차이입니다. 기존 VLA 계열의 대표적인 방법은 Continuous Action을 Discretization하여 Text Token과 비슷하게 Autoregressive하게 예측합니다. 반면 π0는 Continuous Action Distribution 자체를 Flow Matching으로 모델링 합니다.

Conditional Flow Matching

드디어 π0의 핵심인 Flow Matching입니다!!

π0는 Action Chunk의 각 Action에 대응하는 Action Token을 Action Expert에 입력하고 Conditional Flow Matching Loss 로 학습합니다.

\[\mathcal{L}^{\tau}(\theta) = \mathbb{E}_{p(A_t|o_t), q(A_t^{\tau}|A_t)} \left[ \left\| v_{\theta}(A_t^{\tau}, o_t) - u(A_t^{\tau}|A_t) \right\|^2 \right]\]

처음 보면 어려워 보이지만 핵심은 간단합니다.

현재 Noisy Action이 실제 Action으로 가기 위해 어느 방향으로 움직여야 하는지를 나타내는 Vector Field를 학습합니다.

먼저 Gaussian Noise를 샘플링합니다.

\[\epsilon \sim \mathcal{N}(0,I)\]

그리고 실제 Action과 Noise 사이를 다음과 같이 선형으로 Interpolation합니다.

\[A_t^{\tau} = \tau A_t + (1-\tau)\epsilon\]

여기서 \(\tau \in [0,1]\) 입니다.

\(\tau=0\) 이면,

\[A_t^0 = \epsilon\]

즉 완전한 Noise이고,

\(\tau=1\) 이면,

\[A_t^1 = A_t\]

즉 실제 Action이 됩니다.

따라서 Flow Matching은 대략 아래의 경로를 학습하는 것입니다.

\[\text{Noise} \xrightarrow{\text{learned vector field}} \text{Real Action}\]

모델이 맞춰야 하는 Target Vector Field는,

\[u(A_t^{\tau}|A_t) = A_t - \epsilon\]

입니다.

이 식은 위의 Interpolation을 \(\tau\) 로 미분하면 더 직관적입니다.

\[\frac{dA_t^{\tau}}{d\tau} = A_t-\epsilon\]

즉 Noise에서 실제 Action 방향으로 이동하는 속도 벡터를 모델이 예측한다고 생각하면 됩니다.

Robot timestep과 Flow timestep

여기서 \(t\) 와 \(\tau\) 가 같이 등장해서 굉장히 헷갈립니다😵‍💫

  • 아래첨자 \(t\) : 실제 Robot Trajectory의 Timestep
  • 위첨자 \(\tau\) : Flow Matching 과정의 Timestep

완전히 다른 시간입니다.

예를 들어 \(A_{100}^{0.3}\) 이라고 하면, Robot Trajectory의 100번째 시점에서 만들어지는 Action Chunk가 Flow Matching 과정상 0.3 지점에 있다는 뜻입니다.

Bidirectional Attention

Action Expert는 Full Bidirectional Attention Mask를 사용합니다. 즉 Action Chunk 안의 모든 Action Token이 서로 Attend할 수 있습니다.

\[[a_t, a_{t+1}, \cdots, a_{t+49}]\]

를 각각 독립적으로 만드는 것이 아니라 미래 Action끼리 서로 관계를 보면서 전체적인 Action Sequence를 만들 수 있다는 의미입니다.

그리고 학습할 때 \(\tau\) 를 Uniform하게 뽑지 않고 Beta Distribution에서 샘플링하여 낮은 \(\tau\), 즉 Noise가 많은 구간을 더 강조합니다.

Inference

Inference에서는 당연히 정답 Action \(A_t\) 를 알 수 없기 때문에 Random Noise부터 시작합니다.

\[A_t^0 \sim \mathcal{N}(0,I)\]

그리고 학습한 Vector Field를 이용해 \(\tau=0\) 에서 \(\tau=1\) 로 이동합니다. 논문에서는 Forward Euler Integration을 사용합니다.

\[A_t^{\tau+\delta} = A_t^{\tau} + \delta v_{\theta}(A_t^{\tau},o_t)\]

논문에서는 \(\delta=0.1\) 을 사용하기 때문에 총 10번의 Integration Step을 수행합니다.

\[A_t^0 \rightarrow A_t^{0.1} \rightarrow A_t^{0.2} \rightarrow \cdots \rightarrow A_t^{1.0}\]

최종 \(A_t^{1.0}\) 이 Robot에 사용할 Continuous Action Chunk가 됩니다.

그런데 매번 VLM 전체를 10번 다시 계산하면 Real-time Control이 어렵겠죠?? π0는 Observation Prefix의 Attention Key와 Value를 Cache하고, Integration Step마다 Action Token에 해당하는 Suffix만 다시 계산합니다. Transformer에서 흔히 사용하는 KV Cache 를 Flow Matching Inference에도 활용하는 것입니다.

수식에서 주의해서 볼 부분

논문을 공부하면서 개인적으로 가장 헷갈렸던 부분이라 따로 정리해보겠습니다.

1. Vector Field의 부호

Interpolation은,

\[A_t^{\tau}=\tau A_t+(1-\tau)\epsilon\]

이므로 미분하면,

\[\frac{dA_t^{\tau}}{d\tau}=A_t-\epsilon\]

입니다. 따라서 현재 공개된 v4 논문의 Target도,

\[u=A_t-\epsilon\]

으로 적혀 있습니다. Inference가 \(\tau=0\) 의 Noise에서 \(\tau=1\) 의 Action 방향으로 진행되기 때문에 이 부호가 맞습니다.

처음 읽을 때 \(\epsilon-A_t\) 로 방향을 반대로 읽으면 Inference 식과 맞지 않게 되니 Flow가 어느 방향으로 이동하는지를 같이 보시면 헷갈리지 않습니다.

2. Gaussian 표기와 실제 Sampling 식

논문 본문에는 Probability Path를 다음처럼 적습니다.

\[q(A_t^{\tau}|A_t) = \mathcal{N}(\tau A_t,(1-\tau)I)\]

그런데 바로 이어서 실제 Sampling은,

\[A_t^{\tau}=\tau A_t+(1-\tau)\epsilon\]

으로 설명합니다.

여기서 일반적으로 \(\mathcal{N}(\mu,\Sigma)\) 의 두 번째 항을 Covariance로 해석하면 조금 주의가 필요합니다. 아래 Sampling 식에서 Noise에 \((1-\tau)\) 를 곱했기 때문에 실제 Covariance는 \((1-\tau)^2I\) 가 되기 때문입니다.

즉 표준적인 Gaussian 표기법으로 읽는다면,

\[\mathcal{N}(\tau A_t,(1-\tau)^2I)\]

와 Sampling 식이 자연스럽게 대응합니다. 논문이 두 번째 인자를 Scale처럼 표현한 것인지, 단순 표기상의 차이인지는 본문만으로 확정하기 어렵습니다. 그래서 개인적으로는 실제 학습 과정을 설명한 Interpolation 식을 기준으로 이해하는 것이 가장 안전해 보입니다.

이 부분은 논문의 핵심 아이디어를 바꾸는 내용은 아니지만, 수식을 직접 따라가다 보면 한 번쯤 걸릴 수 있는 부분입니다.

기존 Autoregressive VLA와 비교

마지막으로 기존 VLA와 차이를 간단하게 정리해보겠습니다.

구분 기존 Autoregressive VLA π0
Backbone Pre-trained VLM Pre-trained VLM
Action 표현 Action을 Discretization하여 Token처럼 처리 Continuous Action
Action 생성 Autoregressive Prediction Conditional Flow Matching
출력 순차적인 Action Token Action Chunk \(H=50\)
Robot 전용 Weight 일반 VLM Weight에 통합되는 경우가 많음 별도 Action Expert
Action 관계 Causal Prediction Full Bidirectional Attention
강점 VLM과 Action을 동일 Token Framework에서 처리 High-frequency, Dexterous Continuous Control

한 문장으로 표현하면 π0는,

\[\boxed{ \text{VLM의 Semantic Intelligence} + \text{Flow-based Continuous Robot Control} }\]

이라고 볼 수 있을 것 같습니다.

마무리

이번 논문에서 개인적으로 가장 흥미로운 점은 단순히 “VLM에 Action을 붙였다”가 아니라 LLM에서 성공한 Foundation Model의 학습 방식을 Robot Learning으로 상당히 직접적으로 가져왔다는 점입니다.

다양한 데이터로 Pre-training하여 넓은 Physical Capability와 Recovery 능력을 확보하고, 고품질 데이터로 Post-training하여 실제 Task를 잘 수행하도록 만드는 방식입니다. 그리고 기존 Autoregressive VLA에서 Continuous Action을 Discretization하던 방법 대신 Action Expert + Flow Matching을 사용하면서 High-frequency Dexterous Control까지 연결했습니다.

정리하자면,

  • VLM을 통해 Internet-scale Semantic Knowledge 활용
  • Cross-embodiment Robot Data로 Physical Experience 학습
  • Action Expert로 Robotics-specific State/Action 처리
  • Action Chunking으로 미래 Action Sequence 생성
  • Conditional Flow Matching으로 Continuous Action Distribution 모델링
  • Diverse Pre-training + High-quality Post-training으로 Generalization과 Dexterity를 동시에 추구

정도가 π0의 핵심이라고 생각합니다.

최근 VLA 논문들을 읽다 보면 단순히 모델 Architecture만 보는 것보다 어떤 Robot Data를 어떤 Recipe로 학습했는가가 점점 더 중요해지고 있다는 느낌이 듭니다. π0는 이런 Robot Foundation Model 방향을 굉장히 잘 보여주는 논문인 것 같습니다. 긴 글 읽어주셔서 감사합니다😁😁