1.6 Part 1 · Core · foundations

The Transformer

A policy has to read several cameras, a command and the joint state at the same time, and output a sequence of actions. RT-1 (Google, 2022; lesson 3.1), ACT (2023; lesson 1.7), Octo (an open policy, 2024; lesson 3.5), GR00T N1 (NVIDIA, 2025; lesson 3.7) and π0 do this with a transformer. Our running example is π0, a 2024 model from Physical Intelligence; we'll cover it fully in lesson 3.6, and here we take its inputs and inference time as a reference. In this lesson we'll work through attention on six tokens you can move by hand, then count what tokens cost in real policies and see why the control rate is not the same as the inference rate. If you already know the transformer, you can skip this lesson.

≈ 45 minutes1 lab and 5 interactives10 tasks with a goal

TOKENSCut a kitchen frame into patches and work out how much more attention costs

ATTENTIONAim the query at the cup, find an output between objects and keep softmax from collapsing

MASK AND HEADSBuild π0's block mask cell by cell and tune two heads on one scene

ORDERFind commands that can't be told apart without positional encodings

COSTGive π0 a fourth camera and stay within the old attention budget

π0 INFERENCEFind where the 73 ms go and how they give 50 Hz control

Start the lesson ↓

1 Policy inputs

How many inputs a policy has

Take π0, the large 2024 policy from Physical Intelligence, which we'll cover fully in lesson 3.6. At every step the model receives frames from three cameras, a text command like “put the cup in the sink” and the robot's state, and it outputs a chunk of 50 actions. A frame is thousands of pixels, a command is a few words, the state is a short vector of numbers. The inputs differ in nature and size, and a single network has to read them all.

The transformer handles this with one trick: everything is turned into tokens, vectors of the same length. A frame is cut into square patches, a command into words and word pieces, and the state goes through a linear layer and becomes a single token. From then on the network works with a sequence of tokens. The links between them are set up by attention: each token decides which other tokens to look at and what to take from them.

π0 input as tokensabout 870 tokens per step (calculated from the openpi code)

768 frames from three cameras, 256 tokens per 224 px image
up to 48 command
1 robot state
50 actions in the chunk

2 Tokens

How an image becomes tokens

Vision Transformer (ViT, 2021) cuts an image into square patches, flattens each one into a vector and projects it with a linear layer. A patch becomes a token, like a word in a sentence. The smaller the patch, the more tokens there are and the better the network sees details. But attention compares every token with every other, so the number of pairs grows as the square of the number of tokens.

Here is a camera frame of Ada's table, 224 × 224 pixels. Place your bet, then switch the patch size and hover over the frame.

Ada is the home humanoid robot the course labs are built around. More about her in lesson 0.1.

Patch size

49tokens per image

2,401attention pairs,

n^2

Under the hood: ViT and where 196 tokens come from

\begin{gathered} N = \frac{H W}{P^2},\\ z_0 = \bigl[x_{\text{class}};\ x^1 E;\ x^2 E;\ \ldots;\ x^N E\bigr] + E_{\text{pos}},\\ E \in \mathbb{R}^{(P^2 C) \times D}. \end{gathered}

An $H \times W$ image with $C$ channels is cut into $N$ patches of $P \times P$ . Each patch is flattened into a vector of $P^2 C$ numbers and multiplied by the matrix $E$ ; a [class] token is prepended, and a learned position is added to every token. For 224 × 224 and $P = 16$ this gives 14 × 14 = 196 patches plus [class] (our calculation). ViT-B has 12 layers, $D = 768$ , 12 heads and 86M parameters.

A transformer isn't always better than a convolutional network. On ImageNet without strong regularization, ViT is a few points behind ResNet and catches up with the best models only with 14–300 million training images. That's why in RT-1 and ACT the images are encoded by a convolutional network: RT-1 gets 81 tokens from a 300 × 300 frame, and TokenLearner compresses them to 8 tokens per frame. This makes inference 2.4 times faster.

🔬 Lab

Find the cup: attention step by step

Ada is looking for the cup on the table. Each of the six objects is a token. A token has a key: the vector by which it is found. And it has a value: what the token gives back if it is chosen. Here an object's value is its location on the table. The query “where is the cup” is compared with every key by a dot product, softmax turns the scores into weights, and the attention output is the sum of the values with these weights. Everything is computed by hand on two-dimensional vectors that you move yourself. New knobs appear from mission to mission.

Ada's tabletop view; the ring shows the object's weight

Keys and querydrag the tip of the thick arrow

Softmax weights

d_k2 divide by

\sqrt{d_k}

Under the hood: the attention formulas, heads and masks

\begin{gathered} \operatorname{Attention}(Q, K, V) = \operatorname{softmax}\Bigl(\frac{Q K^{\mathsf T}}{\sqrt{d_k}} + M\Bigr) V,\\ Q \in \mathbb{R}^{n \times d_k}, \quad K \in \mathbb{R}^{m \times d_k}, \quad V \in \mathbb{R}^{m \times d_v}. \end{gathered}

Softmax is taken over rows: for each query the weights sum to 1. Why divide by $\sqrt{d_k}$ ? If the components of $q$ and $k$ are independent, with mean 0 and variance 1, then the variance of $q \cdot k$ equals $d_k$ . At large $d_k$ the scores spread out, softmax moves into a region where almost all the weight goes to one token, and the gradients are tiny.

\begin{gathered} \operatorname{MultiHead}(Q, K, V) =\\ \operatorname{Concat}(\text{head}_1, \ldots, \text{head}_h)\, W^O,\\ \text{head}_i = \operatorname{Attention}(Q W_i^Q,\ K W_i^K,\ V W_i^V),\\ W_i^Q, W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}, \quad W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v},\\ W^O \in \mathbb{R}^{h d_v \times d_{\text{model}}}, \quad d_k = d_v = d_{\text{model}} / h. \end{gathered}

Each attention head has its own matrices, so the heads look at different features. A head's width is $d_{\text{model}}/h$ , so all the heads together cost almost as much as a single full-width head. Self-attention: $Q$ , $K$ and $V$ come from the same sequence. Cross-attention: $Q$ comes from one sequence, $K$ and $V$ from another. That's how the action queries in ACT read image features, and how the DiT in GR00T N1 reads VLM features.

The mask $M$ is 0 where token $i$ may look at token $j$ , and $-\infty$ where it may not. A causal mask closes everything to the right of the diagonal: $M_{ij} = -\infty$ for $j > i$ . The block mask of π0 and the prefix-LM in PaliGemma (Google's vision-language model, 2024; lesson 3.2) open a cell if the block number of $j$ is no greater than that of $i$ : within a block attention goes both ways, between blocks only backward. A padding mask closes the columns of empty tokens.

The 2017 block and its feed-forward sublayer:

\begin{gathered} x \leftarrow \operatorname{LayerNorm}\bigl(x + \operatorname{Sublayer}(x)\bigr),\\ \operatorname{FFN}(x) = \max(0,\ x W_1 + b_1)\, W_2 + b_2. \end{gathered}

In ViT and Gemma (Google's language model inside π0; lesson 3.2) the norm comes before the sublayer:

x \leftarrow x + \operatorname{Sublayer}\bigl(\operatorname{Norm}(x)\bigr).

How the lab works. An object has two features: color (red, blue, white) and location (at the front edge, in the middle, at the back of the table). The key of the general head is the color vector plus 0.7 times the location vector, so all six keys are different. The “color” head sees only color, the “location” head only location; their keys have length 1.4. An object's value is its coordinates on the 1 × 1 m table. The query length is at most 3. For $d_k > 2$ , $d_k - 2$ hidden components are added to the two visible ones: random numbers with variance 1, the same on every run (seed 169). Mission thresholds: cup weight above 0.5; the output more than 15 cm from the object centers with no weight above 0.6; the largest weight below 0.9; pairs above 0.8 and the cup above 0.3 in both heads.

3 Positions

Attention doesn't know order

A weighted sum doesn't depend on the order of its terms. If you shuffle the tokens, each one carries its weight along, and the output doesn't change. For a robot this is a problem: “take the cup, then the bowl” and “take the bowl, then the cup” are different commands. So positional encodings are added to the tokens: vectors that depend only on the position number. In the 2017 paper these are sinusoids of different frequencies.

Below is a command of four word tokens. The action token asks them: which object to take first? Rearrange the words and look for commands on which Ada gets it wrong.

Command — drag a word, or click two words to swap them

Add sinusoidal positions

The attention output is 20 cells: the first 4 show how much was taken from each word, the other 16 from which positions; without positions they are empty. The brighter the cell, the larger the number.

0failures found: commands where Ada takes the wrong object first

Similarity of positional encodings

Each cell is the cosine similarity of the encodings of two positions, 16 numbers per position. The diagonal and the cells next to it are the brightest: neighboring positions are similar and distant ones less so, which is why the same word in different places gets different vectors. The frame marks positions 0–3, where the command sits. Hover over a cell.

Under the hood: sinusoids, RoPE and our toy command

\begin{aligned} \mathrm{PE}(\mathit{pos}, 2i) &= \sin\bigl(\mathit{pos} / 10000^{2i/d_{\text{model}}}\bigr),\\ \mathrm{PE}(\mathit{pos}, 2i + 1) &= \cos\bigl(\mathit{pos} / 10000^{2i/d_{\text{model}}}\bigr). \end{aligned}

Each pair of numbers is a sinusoid of its own frequency: from one turn per 6.3 positions down to very slow ones. Neighboring positions get similar vectors, distant ones less similar vectors. The 2017 paper chose sinusoids in the hope that they would extend to sequences longer than those seen in training. Learned positional embeddings gave almost the same result. In ACT the image positions are 2D sinusoids, in Decision Transformer (a transformer for reinforcement learning, 2021) an embedding of the time step, and in Gemma RoPE: it rotates the query and key vectors by an angle that depends on the position.

Our command. A word token is 4 content numbers (one per word) and 16 position numbers; without positions the latter are zero. The action token's query likes objects, the cup slightly more than the bowl: that's how “training” turned out. It also likes the start of the command: its positional part matches the encoding of position 0. The answer is the object with the larger weight. Without positions the cup always wins, and Ada is wrong in all 12 of the 24 commands where the bowl comes before the cup. With positions the earlier object gets the larger weight, and there are no errors in any of the 24 permutations.

4 Cost

What tokens cost

Attention compares every token with every other: double the number of tokens and you get four times as many pairs. But only this term grows quadratically. The Q, K, V projections and the feed-forward block process tokens one at a time, and their cost grows linearly. π0 has about 870 tokens at a model width of 2048, and for now it's these linear parts that take most of a layer's compute. Cameras and resolution are a budget: each new camera adds hundreds of tokens.

The calculator computes the sequence length, a rough cost of one pass and the size of the KV cache. The presets set the inputs of four real policies: RT-1, ACT, Octo and π0. You choose the model width and depth yourself, so all presets are compared on a model of the same size, by default the size of π0.

Preset:

Input

Images Resolution

Slicing

Output

Actions

Model

d_model Layers

867tokens in the sequence

“Patch N” cuts the frame into N × N pixel squares; “stride 32” is a feature map of a convolutional network in which one cell covers 32 pixels, as with ResNet18 in ACT. “64 per frame” and “8 per frame” compress the frame's features to that many tokens, as GR00T N1 and TokenLearner in RT-1 do. d_model is the length of one token's vector, FFN is the layer's feed-forward block, FLOPs is the number of multiplications and additions. The KV cache stores the keys and values of the observation; a KV head is one set of them, and bf16 means 2-byte numbers.

Under the hood: what the calculator computes

attention: $4 n^2 d L$ FLOPs — the product $Q K^{\mathsf T}$ and the mix of the values $V$ ;
projections and FFN: $24 n d^2 L$ FLOPs — the matrices $Q$ , $K$ , $V$ , $O$ and an FFN of width $4d$ ;
KV cache: $2 L H_{kv} d_{\text{head}} T$ numbers, 2 bytes each in bf16.

This is a rough estimate for a standard transformer block: a multiplication and an addition count as two operations, the FFN has width $4 d_{\text{model}}$ , and norms and softmax are not counted. In real models the FFN is sometimes wider, and then the attention share is even smaller. Vaswani et al. compared layers the same way: a self-attention layer costs $O(n^2 d)$ , a recurrent one $O(n d^2)$ , and attention is cheaper as long as $n < d$ . For π0, $n \approx 870$ with $d_{\text{model}} = 2048$ (our calculation).

FlashAttention doesn't remove the quadratic cost: the FLOPs stay $O(n^2 d)$ , but the extra memory is $O(n)$ , and the speedup comes from fewer accesses to GPU memory. We count the KV cache for the prefix, that is, the observation tokens. The Gemma 2B language model inside π0 has 18 layers, one KV head and $d_{\text{head}} = 256$ : about 18 KB per token in bf16, and a prefix of 816 tokens takes about 15 MB (our calculation). The calculator keeps one KV head with $d_{\text{head}} = 256$ for any model width. For π0 the cache takes little memory; what's expensive is time: the encoders and the pass over the observation take 46 ms of the 73.

5 KV cache and chunks

Why 73 ms gives 50 Hz

When a model builds its output over several steps, the keys and values of the tokens it has already processed don't change. They are computed once and stored. This is the KV cache. π0 caches the observation this way, the images and the command: the block mask from the lab (it forbids observation tokens from looking at action tokens) guarantees that the observation doesn't depend on the actions. The model refines the action chunk over 10 flow matching steps (lesson 1.8), and each step reads the ready K and V instead of running the 816 observation tokens through again.

The second idea is chunks. π0 outputs 50 actions at once, and the model is rerun once every 0.5 s. The robot gets a new action every 20 ms, even though the model is called twice a second. Below is π0 inference on an RTX 4090 with three cameras. Click the blocks and the tick tape.

Without KV cache Model off the robot Inference: 73 ms

Control tapeeach cell is a 20 ms tick; click a tick

Under the hood: the inference numbers and how other systems work

π0 has 3.3B parameters: Gemma 2B with the SigLIP image encoder and a 300M action expert. Inference on an RTX 4090 with three cameras: image encoders 14 ms, the pass over the observation 32 ms, 10 flow matching steps 27 ms, 73 ms in total. If the model runs off the robot, the network adds its share: 86 ms. The “no cache” time is our estimate using the same scheme: 14 + 10 × (32 + 2.7) ≈ 361 ms. The KV cache only works at inference: in training and in the first pass everything still has to be computed.

In the real-time chunking paper (Black, Galliker, Levine, 2025), the KV cache prefill of π0 takes 46 ms with a 20 ms control step: more than two ticks pass during one inference. We'll cover how to stitch chunks without pauses or jerks in lesson 1.7.

Other systems are built differently; we'll cover them in Part 3. π0-FAST (a variant of π0 with actions as tokens; lesson 3.6) spends about 750 ms per chunk versus ~100 ms for π0: the 2B model decodes 30–60 action tokens one at a time. GR00T N1 (2.2B parameters, 1.34B of them in the VLM) outputs a chunk of 16 actions in 63.9 ms on an L40; the VLM runs at 10 Hz, the actions at 120 Hz. In Helix (Figure AI, 2025; lesson 3.7) the slow 7B VLM S2 runs at 7–9 Hz, while the fast 80M policy S1 runs at 200 Hz and controls 35 degrees of freedom.

6 Models

How models output actions

They all have a transformer inside, but they output actions in different ways. You can write an action as a string of tokens and generate them one by one, like text. You can output the whole chunk in one decoder pass. Or you can get the chunk from noise in a few steps of diffusion or flow matching. The method determines how many model passes one chunk takes.

🧩 Task

Sort the models by how they output actions

The models on the cards, except π0, are covered in later lessons, and you don't need to know them in advance: the answer is in the “Actions” column of the table below. Drag a card into a bin, or click a card and then a bin. π0-FAST isn't in the table: it's a variant of π0 whose action chunk is written as tokens, like text.

How models turn observations and actions into tokens

A reference for later: the first column gives whose model it is, the year and the lesson where it's covered.

Model	Observations → tokens	Actions	Attention
Decision Transformer, 2021 — a transformer for reinforcement learning	state through a linear layer (a CNN for Atari) plus a return-to-go token	1 token per step, predicted from the state token	GPT, causal mask; K steps make 3K tokens
RT-1, 35M — Google, 2022 (lesson 3.1)	6 frames of 300 × 300 → FiLM-EfficientNet-B3, language via USE → 81 → TokenLearner → 8 per frame, 48 in total	11 dimensions (7 arm, 3 base, 1 mode) × 256 bins	decoder only, 8 layers, 19M
RT-2, 55B / 12B — Google, 2023, a policy on top of a VLM (lesson 3.3)	image and text through the VLM	a string of 8 integers, 256 bins	as in the original VLM
OpenVLA, 7B — an open policy, 2024 (lesson 3.5)	224 × 224, DINOv2 + SigLIP → projector → Llama 2	7 dimensions × 256 bins in place of the 256 least used vocabulary tokens	autoregression
ACT, ~80M — a policy with action chunks, 2023 (lesson 1.7)	4 cameras 480 × 640 → ResNet18 → 300 per camera, plus joints and z — 1202	100 fixed queries → cross-attention → 100 × 14	encoder of 4 layers, decoder of 7, CVAE
Octo, 27M / 93M — an open policy, 2024 (lesson 3.5)	CNN and 16 × 16 patches: 256 (256 × 256) + 64 (wrist 128 × 128), 2 frames; language T5 → 16	readout tokens → diffusion head → chunk	block mask; readout tokens only read
π0, 3.3B — Physical Intelligence, 2024 (lesson 3.6)	SigLIP: 256 per frame (224 px, openpi); state → 1 token	50 tokens through a 300M action expert, flow matching, 10 steps	3 blocks: [images + text], [state], [actions]
GR00T N1, 2.2B — NVIDIA, 2025, a model for humanoids (lesson 3.7)	224 × 224 → SigLIP-2 and pixel shuffle → 64 per frame; features from the 12th LLM layer	16 actions, flow matching, 4 steps	DiT: self-attention over actions, cross-attention to the VLM
Helix, 7B + 80M — Figure AI, 2025, a two-level model for a humanoid (lesson 3.7)	S2 compresses the scene and the command into one latent vector; S1 has a convolutional encoder	35 degrees of freedom at 200 Hz	S1 is an encoder-decoder with cross-attention

7 Check yourself

Eight common misconceptions

8 Takeaways

Lesson takeaways

A policy turns everything into tokens: image patches, command words, the state, actions. The number of tokens is the main lever on cost.
Attention: the query is compared with the keys, softmax gives the weights, and the output is a mix of values. The query length sets the sharpness, and when the weights are spread out the output lands between the options.
Dividing by $\sqrt{d_k}$ keeps the scores in a working range: without it softmax collapses into a single peak.
The mask decides who sees whom. π0's block mask separates the observation from the actions, so the observation can be cached. Heads look at different features, and together they tell apart what each one confuses.
Attention by itself doesn't know order: positional encodings supply it.
Only the attention term is quadratic. The KV cache speeds up inference, and chunks decouple the control rate from the inference rate: 73 ms per inference and 50 Hz on the robot.

Open questions

You want to give a policy a fourth high-resolution camera. What do you calculate before deciding?

How many tokens the camera adds: in PaliGemma a 224 / 448 / 896 px image is 256 / 1024 / 4096 tokens. How much the attention term grows and what its share of the layer cost becomes. How much KV cache it adds and, most importantly, how much time before the first flow step: for π0 that's already 46 ms of 73. Whether inference fits into the interval between chunk launches. And what training will cost: OpenVLA at 384 px instead of 224 trained 3 times longer with no gain in quality. It's often cheaper to compress the tokens: TokenLearner in RT-1 keeps 8 per frame, GR00T N1 keeps 64.

Why can π0 control at 50 Hz if one inference takes 73 ms?

Because the model outputs a chunk of 50 actions, and the robot executes 25 of them at 20 ms each: the model is needed once every 0.5 s, and 73 ms fits into that interval. Inference itself is fast thanks to the KV cache: the observation is processed once, and the 10 flow steps cost 27 ms. Helix and GR00T N1 go further and split a slow VLM (7–10 Hz) from a fast policy (120–200 Hz).

Paper guide

A. Vaswani et al., “Attention Is All You Need”, NeurIPS 2017.

Read closelyThe sections on scaled dot-product attention, multi-head attention and positional encoding, and also “Why Self-Attention” with its table of layer costs.
Safe to skipThe details of training for translation and the BLEU tables: the 65M-parameter base model trained for 12 hours on 8 P100s, and the 213M big model trained for 3.5 days and reached 28.4 BLEU on EN→DE.
Question to keep in mindIn the table, self-attention is cheaper than a recurrent layer as long as $n < d$ . Does this hold for π0, where $n \approx 870$ and $d_{\text{model}} = 2048$ ? What changes if you give π0 448 px images?

Resources

The Illustrated TransformerJay Alammar — the transformer in pictures; there's also The Illustrated GPT-2
The Annotated TransformerHarvard NLP — the 2017 paper line by line, with code
Let's build GPT: from scratch, in code, spelled outAndrej Karpathy — a transformer from scratch, line by line; the code is nanoGPT
Attention in transformers, step-by-step3Blue1Brown, chapter DL6; the intro is DL5, and there's a text version on the site
The Transformer Family Version 2.0Lilian Weng — a survey of variants; the earlier post is Attention? Attention!
Transformer ExplainerAn interactive walkthrough of GPT-2 in your browser; see also LLM Visualization
How caching worksHugging Face — how the KV cache works at inference

➡️

Next: an action chunk in one pass. ACT is a transformer that outputs a chunk of 100 actions in a single decoder pass. Why chunks help, how many actions to take and what to do at the seams — lesson 1.7, “ACT”.

Was the lesson clear?

Found a mistake or inaccuracy? Report it on GitHub

Found this lesson useful? Say thanks — that is how the author knows the course is being read.