Every VLA in Lesson 22.1 solves the same three problems: turn an image into vectors a network can compare, turn an instruction into vectors in the same space, and let the action depend on both. The answers differ in scale by more than four orders of magnitude, from the course’s 110 000-parameter MiniVLA to π0’s 3.3 billion, but they share a structure: encoders produce tokens, a transformer mixes them, and a readout turns some of the mixed tokens into actions.
This lesson opens the MiniVLA trained in Lesson 22.4 and follows one observation through it, tensor by tensor. It then asks two questions that apply to any VLA. Where does the network look, by its attention weights? And what does its action actually depend on, measured by replacing one input at a time? The answers differ, and the difference is the reason the second question exists.
Two instructions, one image. Top row: “push the red puck to the right zone”; bottom row: “push the green puck to the right zone”. The first two columns overlay the action query’s attention over the 64 image cells on the image, for the first and second transformer layer. In layer 1 the attention is spread over the whole image and is identical for both instructions. In layer 2 it concentrates on the named puck, and the predicted action (the blue arrow, from the pusher) turns toward it. The bars show the attention on each word: in layer 2, almost all of it falls on the zone word, “right”, and almost none on the puck’s colour.
The image $o \in [0,1]^{3\times 64\times 64}$ is concatenated with two coordinate channels holding each pixel’s row and column in $[-1, 1]$ (CoordConv, Liu et al., arXiv:1807.03247), so that later layers can read where a feature is as well as what it is. Three strided convolutions reduce it to an $8\times 8$ grid of $D = 64$ channels; each cell, plus a learned position embedding, becomes one of 64 vision tokens. A cell covers 8 × 8 pixels, about 78 mm of table.
The instruction is split into words, mapped to integer ids from a vocabulary of 24 entries, the 22 words of the training instructions plus <pad> and <unk> (an unseen word becomes <unk>), padded to 12, and embedded: 12 language tokens. The pusher’s position is scaled to about $\pm 1$ and mapped linearly to one state token. One learned vector, the action query, is the token whose output becomes the action. Every token also receives a learned embedding of its modality (query, state, language, vision). The sequence has $1 + 1 + 12 + 64 = 78$ tokens.
Each transformer layer updates every token by attention (Vaswani et al., arXiv:1706.03762):
\[\mathrm{Attn}(X) = \mathrm{softmax}\!\Big(\frac{(XW_Q)(XW_K)^\top}{\sqrt{d}} + M\Big)\,XW_V,\]computed in 4 heads of $d = 16$, where $M$ is $-\infty$ on padding positions so that no token reads them. The row of the softmax belonging to the action query is a probability distribution over the 78 tokens: that is the attention map in the figure. Since every token attends to every other, the vision tokens can read the words in layer 1, and the action query can read word-informed vision tokens in layer 2. This is the only place the modalities meet. After two layers the action query’s output passes through a linear map and $\tanh$ to give $(v_x, v_y) \in [-1,1]^2$, scaled to the pusher’s 0.2 m/s.
The same network accepts two other encoders, used in Lessons 22.4 to 22.6:
| Encoder | Tokens | How language enters vision |
|---|---|---|
| patch (the default) | 64 grid cells | only through attention, after encoding |
| keypoint | 16: the expected image position of each of 16 feature channels (spatial softmax, Levine et al., arXiv:1504.00702) and how strongly it fired | before encoding: the mean word embedding scales and shifts the convolution channels (FiLM, Perez et al., arXiv:1709.07871), as RT-1 conditions its image encoder on the instruction |
| oracle | 3: the true puck positions, with a learned identity each | not at all; a diagnostic that reads privileged state, never a policy |
Output (4) quantifies the figure. In layer 1 the action query puts 0.86 of its attention on the image, spread out, with its maximum in cell (1, 5) for both instructions. In layer 2 the image share drops to about 0.4, the pusher’s state token takes 0.49, and the most-attended image cell moves with the instruction: (2, 7), at the red puck, or (3, 1), at the green one. The two predicted actions, $(0.63, -0.81)$ and $(-0.20, 0.94)$, point in opposite directions.
The word bars contradict a naive reading. If attention showed what the network uses, the colour word would matter little, since the query barely attends to “red” or “green” in either layer. Output (5) tests what the action depends on directly. On 400 frames of 40 held-out demonstrations, each input is replaced in turn and the change in the predicted action is measured:
| Input replaced | Mean change of the action (mean action norm 0.57) | Frames where the action changed by more than 0.2 |
|---|---|---|
| another puck named (same zone) | 1.05 | 0.98 |
| another zone named (same puck) | 0.64 | 0.88 |
| another scene’s image | 1.05 | 0.96 |
| another pusher position | 0.90 | 0.95 |
Naming another puck changes the action as much as showing a different scene. The colour word is decisive, although the action query hardly attends to it. The information has to travel another way: the vision tokens can attend to the words in layer 1, and the query reads the vision tokens in layer 2, where its attention lands on the named puck (inference: the routing through the vision tokens is consistent with both measurements but was not isolated here; the challenge below isolates it).
[!recommendation] Attention maps are hypotheses, interventions are evidence An attention weight says how much of one token’s update came from another in one layer. It does not say what the output depends on: information also flows through the other tokens, the residual stream and later layers. Use attention to form a hypothesis about what a VLA uses, then test it by intervening on that input and measuring the change in the action, as in output (5).
Output (2) counts parameters by component. In every variant, the two transformer layers hold most of them (66 944); the patch encoder adds 40 032. Output (3) times one decision on one CPU thread: 0.82 ms for the patch model, 1.11 ms for the keypoint model, 0.21 ms for the oracle, far below the 100 ms control period. For comparison with Lesson 22.1’s table: OpenVLA’s visual encoder alone has about 600 million parameters, and RT-1, at 35 million, runs at 3 Hz on its robot. The MiniVLA’s structure (tokens, attention, a readout query) is the one these systems share; its scale, its absence of pretraining and its single 64 × 64 image are not.
INPUT: the patch MiniVLA of Lesson 22.4 (seed 0); 40 fresh demonstrations
PROCESS: shapes stage by stage; parameters per component; time per decision; attention of the action query; sensitivity to each input
OUTPUT: tables; runs/l22_2/attention.png
```python file=examples/l22_2_architecture.py import os import time from pathlib import Path
import numpy as np import torch
from mjcourse import tabletop as tt from mjcourse.vla import data, language, train from mjcourse.vla.model import MiniVLA, VLAConfig
RUNS = Path(file).resolve().parents[1] / “runs” OUT = RUNS / “l22_2” TEXT = “push the {obj} puck to the {zone} zone”
def inputs(image: np.ndarray, text: str, state: np.ndarray, tokenizer) -> tuple: img = torch.as_tensor(image).permute(2, 0, 1).float().unsqueeze(0) / 255.0 words = torch.as_tensor(tokenizer(text)).unsqueeze(0) s = torch.as_tensor((state[:4] * train.STATE_SCALE)[:2].astype(np.float32)).unsqueeze(0) return img, words, s
def part1(model: MiniVLA, frame: tuple) -> None: img, words, s = frame print(“(1) one observation through the patch MiniVLA (B = 1)”) with torch.no_grad(): x, pad = model.tokens(img, words, s) vision = model.vision(img) fused = model.fusion(x, src_key_padding_mask=pad) action = model(img, words, s) rows = [(“image, RGB in [0, 1]”, img.shape), (“ + 2 coordinate channels (CoordConv)”, (1, 5, 64, 64)), (“ conv stack -> 8 x 8 grid, D = 64 channels”, (1, 64, 8, 8)), (“ vision tokens (one per cell)”, vision.shape), (“instruction, token ids (padded to 12)”, words.shape), (“ language tokens”, (1, 12, 64)), (“pusher position, scaled”, s.shape), (“ state token”, (1, 1, 64)), (“action query (learned)”, (1, 1, 64)), (“sequence: query, state, words, vision”, x.shape), (“ padding mask (True = ignored)”, pad.shape), (“after 2 transformer layers”, fused.shape), (“action: the query’s output -> linear -> tanh”, action.shape)] for label, shape in rows: print(f” {label:<46}{str(tuple(shape)):>18}”) print(f” {int(pad.sum())} of {pad.shape[1]} positions are padding for this {int((words != 0).sum())}-word instruction”)
def part2(vocab: int) -> None: print(“(2) parameters per component”) print(f” {‘component’:<34}” + ““.join(f”{v:>10}” for v in (“patch”, “keypoint”, “oracle”))) groups = {“vision encoder”: (“vision.”,), “word embeddings and positions”: (“words.”, “word_position”), “state, queries, modality embeddings”: (“state.”, “queries”, “modality”), “transformer (2 layers)”: (“fusion.”,), “action head”: (“head.”,)} counts = {} for v in (“patch”, “keypoint”, “oracle”): model = MiniVLA(VLAConfig(vocab=vocab, chunk=1, vision=v, state_dim=2)) counts[v] = {g: sum(p.numel() for n, p in model.named_parameters() if n.startswith(prefixes)) for g, prefixes in groups.items()} counts[v][“total”] = sum(p.numel() for p in model.parameters()) for g in (*groups, “total”): print(f” {g:<34}” + ““.join(f”{counts[v][g]:>10,}” for v in counts))
def part3(vocab: int, frame: tuple) -> None: torch.set_num_threads(1) img, words, s = frame pucks = torch.zeros(1, 6) print(“(3) one decision on one CPU thread (batch 1, median of 300 forward passes)”) for v in (“patch”, “keypoint”, “oracle”): model = MiniVLA(VLAConfig(vocab=vocab, chunk=1, vision=v, state_dim=2)).eval() times = [] with torch.no_grad(): for _ in range(320): t = time.perf_counter() model(img, words, s, pucks) times.append(time.perf_counter() - t) print(f” {v:<10}{1000 * np.median(times[20:]):6.2f} ms”) torch.set_num_threads(min(8, os.cpu_count() or 1))
def attention(model: MiniVLA, frame: tuple) -> list[np.ndarray]: img, words, s = frame out = [] with torch.no_grad(): x, pad = model.tokens(img, words, s) for layer in model.fusion.layers: _, w = layer.self_attn(x, x, x, key_padding_mask=pad, need_weights=True, average_attn_weights=True) out.append(w[0, 0].numpy()) x = layer(x, src_key_padding_mask=pad) return out
def pick_frame(episodes: list[dict]) -> tuple[dict, int, list[str]]: for e in episodes: zone = e[“task”][1] objects = [o for o in tt.OBJECTS if (o, zone) in language.TRAIN_PAIRS][:2] pucks = e[“pucks”][0].reshape(3, 2) d = [pucks[tt.OBJECTS.index(o)] - e[“states”][0][:2] for o in objects] if d[0] @ d[1] < 0: return e, 0, objects raise RuntimeError(“no frame with pucks on opposite sides of the pusher”)
def part4(model: MiniVLA, episode: dict, t: int, objects: list[str], tokenizer) -> dict: zone = episode[“task”][1] print(f”(4) attention of the action query, same image, two instructions naming different pucks”) shown = {} for obj in objects: text = TEXT.format(obj=obj, zone=zone) frame = inputs(episode[“images”][t], text, episode[“states”][t], tokenizer) maps = attention(model, frame) n_words = int((frame[1] != 0).sum()) with torch.no_grad(): action = model(*frame)[0, 0].numpy() for k, m in enumerate(maps, 1): vision, words = m[14:], m[2:2 + n_words] cell = np.unravel_index(int(vision.argmax()), (8, 8)) print(f” ‘{text}’, layer {k}: attention on image {vision.sum():.2f}, words {words.sum():.2f}, “ f”state {m[1]:.2f}, itself {m[0]:.2f}; most-attended image cell (row, col) ({cell[0]}, {cell[1]})”) print(f” predicted action {np.round(action, 2)}”) shown[text] = {“maps”: maps, “n_words”: n_words, “action”: action} return shown
def part5(model: MiniVLA, episodes: list[dict], tokenizer) -> None: rng = np.random.default_rng(0) frames = [(e, t) for e in episodes for t in rng.choice(len(e[“actions”]), min(10, len(e[“actions”])), replace=False)] changes = {k: [] for k in (“another puck named”, “another zone named”, “another scene’s image”, “another pusher position”)} base_norm = [] with torch.no_grad(): for e, t in frames: obj, zone = e[“task”] text = TEXT.format(obj=obj, zone=zone) a0 = model(inputs(e[“images”][t], text, e[“states”][t], tokenizer))[0, 0].numpy() base_norm.append(np.linalg.norm(a0)) other_obj = [o for o in tt.OBJECTS if o != obj and (o, zone) in language.TRAIN_PAIRS] other_zone = [z for z in tt.ZONES if z != zone and (obj, z) in language.TRAIN_PAIRS] o_e, o_t = frames[rng.integers(len(frames))] variants = {“another puck named”: (e[“images”][t], TEXT.format(obj=rng.choice(other_obj), zone=zone), e[“states”][t]), “another zone named”: (e[“images”][t], TEXT.format(obj=obj, zone=rng.choice(other_zone)), e[“states”][t]), “another scene’s image”: (o_e[“images”][o_t], text, e[“states”][t]), “another pusher position”: (e[“images”][t], text, o_e[“states”][o_t])} for k, (img, txt, st) in variants.items(): a = model(inputs(img, txt, st, tokenizer))[0, 0].numpy() changes[k].append(np.linalg.norm(a - a0)) print(f”(5) sensitivity on {len(frames)} frames of 40 held-out demonstrations: mean change of the action when one input “ f”is replaced (mean action norm {np.mean(base_norm):.2f}; actions are in [-1, 1] per axis)”) for k, v in changes.items(): v = np.array(v) print(f” {k:<26} mean change {v.mean():.2f} changed by more than 0.2 in {np.mean(v > 0.2):.2f} of frames”)
def to_pixel(xy: np.ndarray, size: int = 64, fovy: float = 38.0, height: float = 0.88) -> tuple[float, float]: f = 0.5 * size / np.tan(np.radians(fovy) / 2) return -xy[1] * f / height + size / 2, -xy[0] * f / height + size / 2
def figure(episode: dict, t: int, shown: dict) -> None: import matplotlib matplotlib.use(“Agg”) import matplotlib.pyplot as plt
fig, axes = plt.subplots(len(shown), 3, figsize=(10, 3.3 * len(shown)), constrained_layout=True,
gridspec_kw={"width_ratios": [1, 1, 1.4]})
for r, (text, d) in enumerate(shown.items()):
for k in range(2):
ax = axes[r, k]
ax.imshow(episode["images"][t], extent=(0, 64, 64, 0))
grid = d["maps"][k][14:].reshape(8, 8)
ax.imshow(grid / grid.max(), extent=(0, 64, 64, 0), cmap="magma", alpha=0.55, interpolation="nearest")
ax.set_xticks([])
ax.set_yticks([])
ax.set_title(f"layer {k + 1}: image attention ({d['maps'][k][14:].sum():.2f} of total)", fontsize=9)
if k == 1: # the predicted action, from the pusher
u, v = to_pixel(episode["states"][t][:2])
a = d["action"]
ax.annotate("", xy=(u - 14 * a[1], v - 14 * a[0]), xytext=(u, v),
arrowprops={"arrowstyle": "-|>", "color": "#56B4E9", "lw": 2})
words = text.split()
ax = axes[r, 2]
width = 0.38
for k, colour in ((0, "#56B4E9"), (1, "#0072B2")):
ax.bar(np.arange(len(words)) + (k - 0.5) * width, d["maps"][k][2:2 + len(words)], width, color=colour,
label=f"layer {k + 1}")
ax.set_xticks(range(len(words)))
ax.set_xticklabels(words, rotation=30, ha="right", fontsize=8)
ax.set_ylabel("attention weight")
ax.set_title(f'"{text}"', fontsize=9)
ax.legend(frameon=False, fontsize=8)
fig.savefig(OUT / "attention.png", dpi=150)
plt.close(fig)
if name == “main”: OUT.mkdir(parents=True, exist_ok=True) torch.set_num_threads(min(8, os.cpu_count() or 1)) tokenizer = language.default_tokenizer() model = torch.load(RUNS / “vla” / “minivla_patch_s0_pos.pt”, weights_only=False).eval() episodes = [e for e in data.generate(7, 40) if e[“success”]] first = episodes[0] frame = inputs(first[“images”][0], TEXT.format(obj=first[“task”][0], zone=first[“task”][1]), first[“states”][0], tokenizer) part1(model, frame) part2(len(tokenizer.vocab)) part3(len(tokenizer.vocab), frame) episode, t, objects = pick_frame(episodes) shown = part4(model, episode, t, objects, tokenizer) part5(model, episodes, tokenizer) figure(episode, t, shown) print(f”wrote {OUT / ‘attention.png’}”)
The model:
```python file=src/mjcourse/vla/model.py
from __future__ import annotations
from dataclasses import dataclass
import torch
from torch import nn
@dataclass(frozen=True)
class VLAConfig:
vocab: int
width: int = 64
heads: int = 4
layers: int = 2
chunk: int = 4
max_tokens: int = 12
state_dim: int = 4
vision: str = "patch"
keypoints: int = 16
class VisionEncoder(nn.Module):
def __init__(self, width: int):
super().__init__()
self.net = nn.Sequential(
nn.Conv2d(5, 32, 5, stride=2, padding=2), nn.ReLU(), # 32 x 32
nn.Conv2d(32, 32, 3, stride=2, padding=1), nn.ReLU(), # 16 x 16
nn.Conv2d(32, width, 3, stride=2, padding=1), nn.ReLU(), # 8 x 8
nn.Conv2d(width, width, 1))
self.position = nn.Parameter(torch.zeros(1, 64, width))
def forward(self, image: torch.Tensor) -> torch.Tensor:
b, _, h, w = image.shape
ys = torch.linspace(-1, 1, h).view(1, 1, h, 1).expand(b, 1, h, w)
xs = torch.linspace(-1, 1, w).view(1, 1, 1, w).expand(b, 1, h, w)
features = self.net(torch.cat([image, xs, ys], dim=1)) # (B, D, 8, 8)
return features.flatten(2).transpose(1, 2) + self.position # (B, 64, D)
class FiLMKeypointEncoder(nn.Module):
def __init__(self, width: int, keypoints: int):
super().__init__()
self.c1 = nn.Conv2d(5, 32, 5, stride=2, padding=2) # 32 x 32
self.c2 = nn.Conv2d(32, 32, 3, padding=1)
self.c3 = nn.Conv2d(32, keypoints, 3, padding=1)
self.film = nn.Linear(width, 2 * (32 + 32)) # scale and shift for c1 and c2
self.token = nn.Linear(3, width)
self.identity = nn.Parameter(torch.randn(1, keypoints, width) * 0.02)
def forward(self, image: torch.Tensor, language: torch.Tensor) -> torch.Tensor:
b, _, h, w = image.shape
ys = torch.linspace(-1, 1, h).view(1, 1, h, 1).expand(b, 1, h, w)
xs = torch.linspace(-1, 1, w).view(1, 1, 1, w).expand(b, 1, h, w)
g1, b1, g2, b2 = self.film(language).split(32, dim=-1)
x = torch.relu(self.c1(torch.cat([image, xs, ys], 1)) * (1 + g1[..., None, None]) + b1[..., None, None])
x = torch.relu(self.c2(x) * (1 + g2[..., None, None]) + b2[..., None, None])
logits = self.c3(x).flatten(2) # (B, K, 32 * 32)
attention = torch.softmax(logits, dim=-1)
grid = torch.linspace(-1, 1, 32)
gy, gx = torch.meshgrid(grid, grid, indexing="ij")
ex = (attention * gx.flatten()).sum(-1) # expected x of each keypoint
ey = (attention * gy.flatten()).sum(-1)
presence = logits.max(-1).values.tanh()
return self.token(torch.stack([ex, ey, presence], -1)) + self.identity
class OracleObjects(nn.Module):
def __init__(self, width: int):
super().__init__()
self.token = nn.Linear(2, width)
self.identity = nn.Parameter(torch.randn(1, 3, width) * 0.02)
def forward(self, pucks: torch.Tensor) -> torch.Tensor:
return self.token(pucks.view(-1, 3, 2) / 0.3) + self.identity
class MiniVLA(nn.Module):
def __init__(self, cfg: VLAConfig):
super().__init__()
self.cfg = cfg
d = cfg.width
if cfg.vision == "patch":
self.vision = VisionEncoder(d)
elif cfg.vision == "keypoint":
self.vision = FiLMKeypointEncoder(d, cfg.keypoints)
elif cfg.vision == "oracle":
self.vision = OracleObjects(d)
else:
raise ValueError(cfg.vision)
self.words = nn.Embedding(cfg.vocab, d, padding_idx=0)
self.word_position = nn.Parameter(torch.zeros(1, cfg.max_tokens, d))
self.state = nn.Linear(cfg.state_dim, d)
self.queries = nn.Parameter(torch.randn(1, cfg.chunk, d) * 0.02)
self.modality = nn.Parameter(torch.zeros(4, d)) # query, state, language, vision
layer = nn.TransformerEncoderLayer(d, cfg.heads, dim_feedforward=2 * d, dropout=0.0, batch_first=True)
self.fusion = nn.TransformerEncoder(layer, cfg.layers, enable_nested_tensor=False)
self.head = nn.Linear(d, 2)
def tokens(self, image, words, state, pucks=None) -> tuple[torch.Tensor, torch.Tensor]:
b = image.shape[0]
q = self.queries.expand(b, -1, -1) + self.modality[0]
s = self.state(state).unsqueeze(1) + self.modality[1]
embedded = self.words(words)
w = embedded + self.word_position + self.modality[2]
if self.cfg.vision == "keypoint":
known = (words != 0).float().unsqueeze(-1)
summary = (embedded * known).sum(1) / known.sum(1).clamp(min=1) # mean word embedding
v = self.vision(image, summary) + self.modality[3]
elif self.cfg.vision == "oracle":
v = self.vision(pucks) + self.modality[3]
else:
v = self.vision(image) + self.modality[3]
pad = torch.cat([torch.zeros(b, q.shape[1] + 1, dtype=torch.bool), words == 0,
torch.zeros(b, v.shape[1], dtype=torch.bool)], dim=1)
return torch.cat([q, s, w, v], dim=1), pad
def forward(self, image, words, state, pucks=None) -> torch.Tensor:
x, pad = self.tokens(image, words, state, pucks)
fused = self.fusion(x, src_key_padding_mask=pad)
return torch.tanh(self.head(fused[:, :self.cfg.chunk])) # (B, CHUNK, 2)
Output (MuJoCo 3.14.0, PyTorch 2.14 CPU; about 2 minutes, of which most is generating the demonstrations):
(1) one observation through the patch MiniVLA (B = 1)
image, RGB in [0, 1] (1, 3, 64, 64)
+ 2 coordinate channels (CoordConv) (1, 5, 64, 64)
conv stack -> 8 x 8 grid, D = 64 channels (1, 64, 8, 8)
vision tokens (one per cell) (1, 64, 64)
instruction, token ids (padded to 12) (1, 12)
language tokens (1, 12, 64)
pusher position, scaled (1, 2)
state token (1, 1, 64)
action query (learned) (1, 1, 64)
sequence: query, state, words, vision (1, 78, 64)
padding mask (True = ignored) (1, 78)
after 2 transformer layers (1, 78, 64)
action: the query's output -> linear -> tanh (1, 1, 2)
4 of 78 positions are padding for this 8-word instruction
(2) parameters per component
component patch keypoint oracle
vision encoder 40,032 27,504 384
word embeddings and positions 2,304 2,304 2,304
state, queries, modality embeddings 512 512 512
transformer (2 layers) 66,944 66,944 66,944
action head 130 130 130
total 109,922 97,394 70,274
(3) one decision on one CPU thread (batch 1, median of 300 forward passes)
patch 0.82 ms
keypoint 1.11 ms
oracle 0.21 ms
(4) attention of the action query, same image, two instructions naming different pucks
'push the red puck to the right zone', layer 1: attention on image 0.86, words 0.11, state 0.02, itself 0.01; most-attended image cell (row, col) (1, 5)
'push the red puck to the right zone', layer 2: attention on image 0.42, words 0.09, state 0.49, itself 0.00; most-attended image cell (row, col) (2, 7)
predicted action [ 0.63 -0.81]
'push the green puck to the right zone', layer 1: attention on image 0.86, words 0.11, state 0.02, itself 0.01; most-attended image cell (row, col) (1, 5)
'push the green puck to the right zone', layer 2: attention on image 0.39, words 0.12, state 0.49, itself 0.00; most-attended image cell (row, col) (3, 1)
predicted action [-0.2 0.94]
(5) sensitivity on 400 frames of 40 held-out demonstrations: mean change of the action when one input is replaced (mean action norm 0.57; actions are in [-1, 1] per axis)
another puck named mean change 1.05 changed by more than 0.2 in 0.98 of frames
another zone named mean change 0.64 changed by more than 0.2 in 0.88 of frames
another scene's image mean change 1.05 changed by more than 0.2 in 0.96 of frames
another pusher position mean change 0.90 changed by more than 0.2 in 0.95 of frames
A first example that showed nothing. The first version of output (4) used the sixth frame of the first held-out episode. Attention moved to the named puck, but the two predicted actions were almost equal, $(-0.84, -0.59)$ and $(-0.89, -0.56)$: both pucks lay below the pusher, so both instructions called for nearly the same motion at that moment. The script now picks, by a fixed rule, the first frame in which the two pucks lie in directions from the pusher more than 90 degrees apart.
Padding that leaks. If the padding mask is not passed to every layer, the query attends to <pad> embeddings, and the action changes with the instruction’s length. Check: append a padding token by hand and confirm the output is unchanged to float precision.
Coordinates lost. Without the coordinate channels, a convolutional encoder must infer position from image borders through several layers; with them, every feature carries its position from the first layer. Removing them is a one-line ablation worth running (exercise).
Remove the two coordinate channels from VisionEncoder (5 input channels become 3), retrain with examples/l22_4_train_minivla.py --vision patch --seed 0 and compare in-distribution success with the CoordConv model’s 0.95. Predict the direction of the change before running.
Isolate the route the colour word takes. In layer 1 only, mask the vision tokens’ attention to the words (an attention mask of $-\infty$ from vision rows to word columns), and repeat output (5) for “another puck named”. If the change of the action collapses, the colour reaches the action through the vision tokens; if it does not, it travels through the query directly in layer 2.
[!research] What scales and what does not RT-1 (arXiv:2212.06817) conditions its image encoder on language with FiLM and compresses each image to a few tokens with TokenLearner before a transformer; RT-2 (arXiv:2307.15818) and OpenVLA (arXiv:2406.09246) feed image tokens and words to a pretrained language model and read actions out as text tokens; Octo (arXiv:2405.12213) reads its transformer through readout tokens into a diffusion head; π0 (arXiv:2410.24164) attaches a separate action expert trained with flow matching. The MiniVLA’s readout query follows the same pattern as ACT’s queries (arXiv:2304.13705). None of these papers’ results transfer to a 110 000-parameter model trained from scratch on one scene; what transfers is the vocabulary (tokens, fusion, readout) and the methods of this lesson for asking what a given network uses.
{"id": "22.2-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "How many tokens enter the MiniVLA's transformer, and from where?",
"options": ["64, all from the image", "78: 1 action query, 1 state, 12 words (padded), 64 image cells", "12 words only", "2: an image token and a language token"],
"answer": 1,
"explain": "<p>1 + 1 + 12 + 64 = 78; padding positions are masked.</p>"},
{"kind": "mcq", "q": "The action query hardly attends to the colour word, yet naming another puck changes the action by 1.05 on average. What follows?",
"options": ["The colour word is unused", "Attention weights do not show what the output depends on; the colour can reach the action through other tokens, such as the vision tokens in layer 1", "The sensitivity test is wrong", "The model ignores language"],
"answer": 1,
"explain": "<p>Attention is a hypothesis about information flow; the intervention is the evidence.</p>"},
{"kind": "mcq", "q": "What do the two coordinate channels (CoordConv) give the vision encoder?",
"options": ["Colour constancy", "Each pixel's position as an input, so features can carry where they are", "Depth", "Language conditioning"],
"answer": 1,
"explain": "<p>Convolutions are translation-equivariant; the extra channels let them depend on position.</p>"},
{"kind": "mcq", "q": "Where in the MiniVLA do vision and language first interact?",
"options": ["In the vision encoder", "In the first transformer layer's self-attention", "In the action head", "In the tokenizer"],
"answer": 1,
"explain": "<p>For the patch encoder, attention is the only fusion; the keypoint encoder also uses FiLM before attention.</p>"}
]}
Lesson 22.3 fixes the network and changes what its output means: torques, joint targets, velocities or Cartesian displacements, and what each does to cloning and robustness.