Robot Foundation Models

π0: A Vision-Language-Action Flow Model for General Robot Control

Successor to Octo model.

Links:

• https://www.physicalintelligence.company/blog/pi0
• https://www.youtube.com/live/ELUMFpJCUS0?t=16866s kevin black motivating the architectural design (coming from Octo) starting at 04:41:06
• Me trying to explain how the architecture works: https://www.youtube.com/watch?v=NVS-7VJMA5c

Two main contributions:

1. Applying VLM to VLA via flow matching
   1. Leverages large-scale internet pretraining from VLM
2. Data recipe

” averaging over 10 trials per task”

• This is how many trials they do to get success rate

Model Architecture

It’s a MoE-style architecture composed of 2 experts:

1. Expert 1 (VLM): PaliGemma (3 billion params) pre-trained from internet
2. Expert 2(large transformer): Gemma action expert (300 million params) from scratch
   • The Gemma expert is a dedicated transformer stack for actions

The two experts talk to each other via Blockwise causal attention.

• This is a really important detail that is not shown in the paper

• There are 18 transformer blocks (i.e. depth in the paper)
• Width = the model’s hidden size (a.k.a. embedding dim, model dim).
  • For the VLM expert (PaliGemma backbone): width = 2048
  • For the Action expert: width = 1024
• Note that the width doesn’t have to match between the two experts, just during attention, the head dim needs to be the same
  • Both experts map into the same attention head space ($18$ heads $\times 256$ hidden dim = $4608$).
  • Concatenate tokens → do one global attention.
  • Split results → project back into each expert’s private width (2048 vs 1024).
  • Enables cross-attention between VLM and Action tokens, while keeping parameters separate.

Why use separate expert as opposed to just introducing new tokens to the VLM?

That makes it a giant decoder, and is what Kevin black actually originally tried for the pi0 model. However, convergence is really slow, and distribution shift. it will make the model super confused ( explained at ~5:01:31 of the talk)

Idea: use a whole set of different weights for the action expert, which is trained from scratch.

Also another thing kevin pointed out: leveraging pre-training is super-duper important!!

Flow Matching

At training time, the Flow Matching loss to train the policy is given by $L_{\tau }(\theta )={\mathbb{E}}_{p(A_t\mid o_t),q(A_t^{\tau }\mid A_t)}{\Vert v_{\theta }(A_t^{\tau },o_t)-u(A_t^{\tau }\mid A_t)\Vert }^2$

Where

• $u(A_t^{\tau }\mid A_t)=ϵ-A_t$ (Note: this should be $A_t-ϵ$?)
• $A_t^{\tau }=\tau A_t+(1-\tau )ϵ$

Potential source of confusion

Notice that $v_{\theta }(A_t^{\tau },o_t)$ is always learning to predict $ϵ-A_t$, even though it is conditioned on $A_t^{\tau }$. You might think really it should be learning $\tau (ϵ-A_t)$, but that would be learning the distance vector - We are trying to learn the velocity field, which stays constant through (the first derivative).

• This multiplication by $\tau$ will be done at inference time to control “step size”

What's the point of \tau?

Without $\tau$, where you just start from $ϵ$ and directly predict $A_t$, that’s essentially a denoising autoencoder view.

$\tau$ allows you to better capture multi-modal behavior, else you end up with mode-collapse, think about this scenario:

Background on flow matching

ChatGPT conversation: https://chatgpt.com/share/68c5b5fa-7db4-8002-9708-ef4e953533f9

The idea: we want to turn noise $ϵ\sim \mathcal{N}(0,I)$ into a data sample (here, an action $A_t$). We can describe this transformation as continuous process over time $\tau \in [0,1]$, i.e. an ODE:

$$\frac{d}{d\tau }{X}^{\tau }=v(X^{\tau },\tau ),$$

with initial condition $X^0=ϵ$ and at the end, $X^1\approx A_t$. Flow matching tries to learn this vector field $v(\cdot )$.

To train such a model, we choose reference path between $ϵ$ and $A_t$ and differentiate over it. The simplest path between $ϵ$ and $A_t$ is a straight line (i.e. straight line interpolation): $A_t^{\tau }=(1-\tau )ϵ+\tau A_t,\tau \in [0,1]$

• At $\tau =0$, you are at pure noise, and at $\tau =1$, you’re at the action

Taking the derivative with regards to $\tau$, we see that: $\frac{d}{d\tau }{A}_t^{\tau }=\frac{d}{d\tau }((1-\tau )ϵ+\tau A_t)=A_t-ϵ$

• The derivative of a straight line is a constant slope, so we just need to learn this constant!

We learn to predict this gradient (a constant) so that at inference time, we learn this mapping for any arbitrary $A_t^{\tau }$ $\frac{d}{d\tau }{A}_t^{\tau }=v_{\theta }(A_t^{\tau },o_t)=A_t-ϵ$

At inference time, we start with random noise $A_0^t\sim \mathcal{N}(0,I)$ and integrate the learned vector field from $\tau =0$ to $\tau =1$, and use forward Euler integration rule: $A_t^{\tau +\delta }=A_t^{\tau }+\delta v_{\theta }(A_t^{\tau },o_t)$

• where $\delta$ is the integration size ($\delta =0.1$ in paper)

Why is 10 steps better than 1 step?

Because at the end of the day, we are trying to learn multi-modal distributions.

• ChatGPT answer: If you take 10 smaller steps, each step only needs to be locally correct. Integration keeps pulling you back onto the line. So error doesn’t explode; it averages out