π0.5 : a Vision-Language-Action Model with Open-World Generalization
연구실 밖의 현실 세계 상황에서는 데이터 스케일을 늘리는 것만으로는 부족함. 새로운 물체를 어떻게 잡을지와 같은 low-level physical behavior부터 어떤 물체가 target인지 판단하는 것과 같은 high-level scene semantics까지 여러 abstraction level에서 일반화할 수 있어야 한다. 따라서 이러한 폭넓은 지식을 제공할 수 있는 training recipe를 설계할 필요가 있다.
중요한 문제 제기 : 유연한 open-world generalization을 위해서 로봇 학습 시스템의 training recipe를 어떻게 설계해야 하는가?
가설 : generalizable robot learning system은 경험과 다양한 정보로부터의 지식을 transfer 할 수 있어야한다.
저자들은 이러한 가설을 기반으로 다양한 modality와 task를 동일한 VLA framework에서 학습할 수 있도록 heterogeneous co-training framework을 설계하고 이를 통해 unseen environment에서도 broad generalization이 가능한 $\pi_{0.5}$를 제안.
모델 구조
Contribution •Heterogeneous co-training enables broad generalization •Unified high-level reasoning + low-level action generation •Long-horizon manipulation in unseen real homes
Limitation •Partial observability & manipulation failures •Unstable high-level subtask inference •Limited context / memory & prompt complexity
Future Works •Heterogeneous co-training enables broad generalization •Unified high-level reasoning + low-level action generation •Long-horizon manipulation in unseen real homes
VLA의 학습 목적은 Observation + instruction을 조건으로 action chunk를 예측하는 imitation learning의 log-likelihood의 최대화이다.
\(\max_\theta \mathbb{E}_{(\mathbf{a}_{t:t+H}, \mathbf{o}_t,l)\sim\mathcal{D}}\log(\pi_\theta(\mathbf{a}_{t:t+H}\vert \mathbf{o}_t,l))\) 이 때, observation에는 이미지 뿐만이 아니라 joint position 같은 로봇 내부 상태도 들어간다.
\[\mathbf{o}_t=\{I_t^1, \ldots, I_t^n,q_t\}\]VLA는 현대적인 LLM/VLM 구조를 따르며, modality-specific tokenizer를 통해 입출력을 discrete(hard) 또는 continous(soft) token representation으로 변환한다. 이후에 대규모 autoregressive transformer backbone이 입력 토큰을 조건으로 출력 토큰을 순차적으로 예측하도록 학습된다.
모델의 weight는 pretrained VLM으로 초기화된다.
Imitation Learning은 간단한 next-token predcition 문제로 생각할 수 있다.
\[(\mathbf{o}_t,l,\mathbf{a}_t) \rightarrow[\text{obs tokens, language tokens, action tokens}]\]이를 통해 현대 ML의 scalable training tool을 그대로 활용할 수 있다.
이미지와 텍스트 입력의 tokenizer는 moder VLM에서 사용하는 방식을 따른다.
Action의 경우에 pretraining에서는 기존 연구의 compression-based tokenization approach를 사용한다.
Post-Training 단계에서는 $\pi_0$처럼 Flow matching을 이용해 action distibution을 표현한다.
이전 flow matching step의 Partially denoised action $\mathbf{a}$를 action에 해당하는 action token들이 입력받아서 Flow Matching Vector Field를 출력한다. 그리고 이 과정을 여러번의 flow matching step에 걸쳐 반복하면서 최종 action에 가까워진다.
Action Token을 처리할때는 LLM/VLM Backbone의 weight를 그대로 사용하는게 아니라 별도의 weight set을 사용한다. 저자는 이를 Action Expert라고 표현한다. Mixture of Expert 구조와 유사하게 토큰 종류에 따라 다른 parameter set을 사용한다.
굳이 action expert를 따로 두는 이유는 LLM은 비전/언어 등을 포함한 범용 representation 처리를 해야하지만, Action Expert는 flow matching based action generation이라는 특정한 역할만 담당하며 LLM보다 훨씬 작게 구성할 수 있다.
\[\begin{align} (\mathbf{o}_t,l)&\rightarrow \text{LLM/VLM Backbone} \\ \mathbf{a}^\tau &\rightarrow \text{Action Expert} \\ & \Downarrow \\ v_\theta(\mathbf{a}^\tau,&\tau\vert\mathbf{o}_t,l) \end{align}\] 모델 가중치는 웹데이터로 훈련된 VLM으로 초기화된다.
훈련은 pre-training, post-training 두 단계로 나뉜다.
Pre-training 단계는 다양한 로보틱 task에 모델이 적응하도록 하고, post-training 단계에서는 모델을 mobile manipulation에 특화시키고 효율적인 test-time inference를 위한 매커니즘을 갖추도록 설계되었다.
Pre-training 단계에서는 모든 task가 simple, scalable, efficient training을 위해 discrete token으로 표현된다. Post-training 단계에서는 모델에 action expert를 적용하여 더 세밀한 단위로 동작을 표현하고 실시간 제어를 위한 효율적인 연산 inference를 가능하게 한다.
Inference 단계에서는 모델이 수행해야할 고수준의 subtask를 생성하고 이 subtask에 따른 저수준의 action을 action expert를 통해 예측한다.
모델은 Action Chunk Distribution과 Tokenized Text Output을 유연하게 표현할 수 있다. Tokenized Text Output은 co-training task와 계층적 추론 과정 중에서 고수준의 subtask prediction을 출력하기 위해서 사용할 수 있다.
모델은 observation $\mathbf{o}t$와 task prompt $\mathcal{l}$을 입력받아 action chunk $\mathbf{a}{t:t+h}$와 textual output $\hat{\mathcal{l}}$을 모델링한다. $\hat{\mathcal{l}}$은 고수준 subtask 예측이나 vision-language prompt에 대답하는데 사용된다.
\[\pi_\theta(\mathbf{a}_{t:t+H},\hat{\mathcal{l}} \vert \mathbf{o}_t,\mathcal{l}) = \underbrace{\pi_\theta(\mathbf{a}_{t:t+H} \vert \mathbf{o}_t,\hat{\mathcal{l}})}_{Low-level} \underbrace{\pi_\theta(\hat{\mathcal{l}} \vert \mathbf{o}_t,\mathcal{l})}_{High-level}\]이때 하나의 모델이 high-level, low-level inference를 모두 수행한다. High-level inference는 뒤의 항을 표현하고, low-level inference는 앞의 항을 표현한다. 이 때 action distribution은 original task prompt가 아닌 $\hat{\mathcal{l}}$에 의존한다.
Transformer는 text token, image patch, intermediate denoising action으로 구성된 멀티모달 입력 $x_{1:N}$ 값을 받아 멀티모달 출력 $y_{1:N}=f(x_{1:N},A(x_{1:N}),\rho(x_{1:N}))$을 뱉는다.
Token type에 따라 서로 다른 encoder와 expert weight를 사용할 수 있으며, 이 타입은 $\rho(x_i)$로 표현한다.
일반적인 LLM과 달리 image patch, textual prompt, continuous action token에는 bidirectional attention을 사용한다.
출력 $y$는 다음처럼 나눌 수 있다.
\(y\rightarrow[y_{1:M}^l, y^a_{1:H}]\) $y^l$은 $\hat{l}$을 샘플링 하기 위한 text logits, $y^a$는 Action Expert가 생성하는 action output token이며, 이후 linear mapping을 통해 continous action을 얻는 데 사용한다.
$\pi_{0.5}$는 Flow matching을 이용하여 Continuous Action을 예측한다.
\(\mathbf a_{t:t+H}^{\tau,\omega} = \tau \mathbf a_{t:t+H} + (1-\tau)\omega, \qquad \omega\sim\mathcal N(0,I)\) 모델은 이 입력으로부터 flow vector field $\omega-\mathbf{a}$를 예측하도록 학습된다.
Discrete / Continuous Representation에는 trade-off가 존재한다. Action Chunk는 압축하여 discrete token으로 표현할 수 있기 때문에 VLA training이 빠르다. 하지만 inference에서는 autoregressive decoding이 필요하기 때문에 real-time inference에 불리하다.
Flow Matching은 Continuous action을 생성하며, iterative flow integration을 통해 action을 생성할 수 있어 최종 모델의 빠른 inference에 사용된다.
즉 training에서는 discrete action의 장점을 활용하고 inference에서는 flow matching을 사용할 수 있는 모델이다. 이때 attention matrix를 이용해서 두 action representation이 서로 참조하지 못하도록 한다.
Inference에서는 text output $\hat{l}$을 autoregressive하게 생성한 뒤, 이를 condition으로 10-step denoising을 수행하여 action chunk를 생성한다.
첫 단계에서 robot + non-robot 데이터를 혼합하여 standard autoregressive next-token prediction으로 pre-training한다.
prediction 대상은 text/obj location/FAST-encoded action token
주요 training data
특히나 HL 학습을 통해 하나의 모델이 high-level command -> subtask -> action 형태의 high-level + low-level policy역할을 모두 수행할 수 있게 한다.
280k step의 discrete-token pre-training 이후 80k step post-training을 수행
목적은 두가지
Post-training에서는
\[\mathcal L_{\mathrm{next-token}} + 10\,\mathcal L_{\mathrm{flow}}\]으로 text prediction capability를 유지하면서 Action Expert를 학습한다.
데이터는 주로 successful MM/ME episodes + HL + WD를 사용하며, 추가로 VI (Verbal Instruction demonstrations)를 도입한다.
VI는 expert가 학습된 low-level policy에 적절한 subtask command를 단계별로 제공한 demonstration으로, 모델의 high-level subtask prediction을 개선하기 위해 사용된다.
