pi 0.5: a Vision-Language-Action Model with Open-World Generalization

Paper Link: https://www.pi.website/download/pi05.pdf

Vision Language Action (VLA) model

VLA total params = Vision Language Model (VLM) params + action head params

Vision Language Action (VLA) models are initialized with the weights of the pre-trained VLM

Training §

pi 0.5 paper goes very into detail about the training mixtures used to train the model which helps improve its generalizability, which I am not going to go into detail here

VLAs are trained via imitation learning on robot demonstrations dataset D, by maximizing log-likelihood of an action $a_t$ given an observation $o_t$ and natural language task $l$

$$ma{x}_{\theta }{\mathbb{E}}_{(a_{t:t+H,o_t,l})\text{textasciitilde}D}log({\pi }_{\theta }(a_{t:t+H}|o_t,l))$$

The observation typically contains one or more images $I_t^1,\dots ,I_t^n$

pi 0.5 is trained in 2 stages:

1. pre-training stage to adapt the model to diverse robotic tasks
2. post-training stage is intended to specialize it to mobile manipulation

post-training is when they add the action expert

Architecture §

${\pi }_{0.5}$ architecture can represent both action chunk distributions and tokenized text outputs

For subtask inference it’s written as

${\pi }_{\theta }(a_{t:t+H},\hat{l}|o_t,l)$

$o_t=[I_t^1,\dots ,I_t^n,q_t]$

• $l$ - overall task prompt
• $\hat{l}$ - high level subtask / answer to vla prompt in web-data
• $a_{t:t+H}$ - predicted action chunk

We can decompose this as

$${\pi }_{\theta }(a_{t:t+H},\hat{l}|o_t,l)={\pi }_{\theta }(a_{t:t+H}|o_t,\hat{l}){\pi }_{\theta }(\hat{l}|o_t,l)$$

Action distribution does NOT depend on $l$ (overall prompt) but depends on the high level subtask $\hat{l}$

Model takes in N multimodal input tokens ($x_{1:N}$), token implies discretized and continuous inputs

We produce a sequence of multimodal outputs $y_{1:N}$, which we write as $y_{1:N}=f(x_{1:N},A(x_{1:N}),p(x_{1:N}))$

• $p(\cdot )$ - token type
• $A(\cdot )$ - Attention Mask, tells model which tokens are allowed to attend to which ones
  • Attention Mask (Detail)
  • Prefix block (P)
    • P tokens attend bidirectionally to each other (images/prompt/state can interact freely)

The model outputs both text and actions the output of $f$ is split into text token logits and action output tokens ($y_{1:M}^l,y_{1:H}^a$)

First M are text token logits that can be used to sample $\hat{l}$, later H tokens are produced by a separate action expert

$M+H\le N$

Discrete and Continuous Action Representations §