Two halves, one policy
You have already built both halves of a vision-language-action policy. In chapter 1.7
you built the data: a multi-task pile of (instruction_tokens, image_features, state) -> action examples, with a from-scratch word-level tokenizer and a frozen, random-init tiny
CNN standing in for a vision backbone. In chapter 1.5 you built the action head: flow
matching — learn the velocity of the straight noise→data line, then sample an action by
integrating an ODE. This chapter fuses them into one language-conditioned policy and trains
it. Nothing here comes from a model zoo: the "VLM" is a token embedding plus a few lines of
from-scratch attention, and the head is the ch1.5 velocity field, now conditioned on what the
attention produced.
That word — fuses — is the whole chapter. A VLA has to take three very different things
(words, pixels-as-features, numbers) and turn them into one vector an action head can
condition on. Open vla.py. It has seven regions: setup, data (consume ch1.7's
.npz), vision+language (rebuild ch1.7's frozen encoder + tokenizer for eval),
model (the tiny VLM + flow head), train, eval (with ch1.6 error bars), and
report.
The fusion backbone, from scratch
A transformer block is not magic. It is three nn.Linear projections (query, key, value),
one scaled dot-product, a softmax, and an output projection — plus a per-token MLP, wrapped
in pre-norm residuals. Here it is, the entire attention mechanism:
Three linear projections (query, key, value), one scaled dot-product, a softmax over the sequence, and a weighted sum of values — the row that lets every token read every other — and here it is in code:
# The tiny VLM + flow head, from scratch. No transformers, no einops — a transformer# block is a from-scratch multi-head attention (Q,K,V projections, a scaled dot-product# softmax, an output projection) plus an MLP, with pre-norm residuals. That is the whole# "backbone"; the lesson is the FUSION, not a big architecture.def sinusoidal_embed(t: torch.Tensor, dim: int) -> torch.Tensor: """Continuous flow time (B,) -> (B, dim) sinusoidal features (ch1.5 / ch1.4).""" half = dim // 2 freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=t.device) / half) ang = t.float()[:, None] * freqs[None] return torch.cat([ang.sin(), ang.cos()], dim=1) class Block(nn.Module): """One pre-norm transformer block: self-attention that lets vision, state, and each word token exchange information, then a per-token MLP. Built from nn.Linear only.""" def __init__(self, dim: int, heads: int) -> None: super().__init__() self.heads = heads self.ln1, self.ln2 = nn.LayerNorm(dim), nn.LayerNorm(dim) self.qkv = nn.Linear(dim, 3 * dim) self.proj = nn.Linear(dim, dim) self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)) self.last_attn = None # CLS-token attention over the sequence, for the rerun viz def forward(self, x: torch.Tensor, key_pad: torch.Tensor) -> torch.Tensor: B, L, dim = x.shape h, hd = self.heads, dim // self.heads qkv = self.qkv(self.ln1(x)).reshape(B, L, 3, h, hd).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # each (B, h, L, hd) scores = (q @ k.transpose(-2, -1)) / math.sqrt(hd) # (B, h, L, L) scores = scores.masked_fill(key_pad[:, None, None, :], float("-inf")) # ignore <pad> keys attn = scores.softmax(dim=-1) self.last_attn = attn[:, :, 0, :].mean(1).detach() # (B, L): how CLS attends to each input x = x + self.proj((attn @ v).transpose(1, 2).reshape(B, L, dim)) return x + self.mlp(self.ln2(x)) class TinyVLA(nn.Module): """Fuse instruction + vision + state -> one conditioning vector, then predict the flow velocity of the action conditioned on it. The sequence is [CLS, vision, state, tok_0..tok_15]; the CLS output (after the blocks) is the fused representation the flow head sees. blind=True zeros the vision input (Break It).""" def __init__(self, vocab_size: int, feat_dim: int, dim: int, layers: int, heads: int, hidden: int, stats: dict) -> None: super().__init__() self.blind = BLIND self.tok_embed = nn.Embedding(vocab_size, dim, padding_idx=PAD_ID) self.vision_proj = nn.Linear(feat_dim, dim) self.state_proj = nn.Linear(STATE_DIM, dim) self.cls = nn.Parameter(torch.zeros(1, 1, dim)) self.pos = nn.Parameter(0.02 * torch.randn(1, 3 + MAX_TOKENS, dim)) self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(layers)]) self.norm = nn.LayerNorm(dim) self.vel = nn.Sequential( nn.Linear(ACT_DIM + TIME_DIM + dim, hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.SiLU(), nn.Linear(hidden, ACT_DIM), ) for name, value in stats.items(): self.register_buffer(name, torch.from_numpy(value)) def fuse(self, tokens: torch.Tensor, img_feat: torch.Tensor, state: torch.Tensor) -> torch.Tensor: B = tokens.shape[0] feat = (img_feat - self.feat_mean) / self.feat_std if self.blind: # Break It: the policy gets NO vision — must solve from words + state feat = torch.zeros_like(feat) st = (2.0 * (state - self.state_min) / self.state_range - 1.0).clamp(-1.0, 1.0) seq = torch.cat([ self.cls.expand(B, -1, -1), # a learned query that reads out the fusion self.vision_proj(feat)[:, None], # the frozen image feature, as one token self.state_proj(st)[:, None], # the state, as one token self.tok_embed(tokens), # the instruction, one token per word id ], dim=1) + self.pos key_pad = torch.zeros(B, 3 + MAX_TOKENS, dtype=torch.bool, device=tokens.device) key_pad[:, 3:] = tokens == PAD_ID # attention ignores padded instruction slots for blk in self.blocks: seq = blk(seq, key_pad) return self.norm(seq[:, 0]) # the CLS row: the fused conditioning vector def velocity(self, x_t: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: return self.vel(torch.cat([x_t, sinusoidal_embed(t * TIME_SCALE, TIME_DIM), cond], dim=1)) def flow_loss(model: TinyVLA, x0: torch.Tensor, mask: torch.Tensor, tokens: torch.Tensor, img: torch.Tensor, state: torch.Tensor) -> torch.Tensor: """ch1.5's conditional flow-matching loss, now conditioned on the fused VLM vector and MASKED to each embodiment's real action dims (pusht's padded 2:6 never train).""" t = torch.rand(len(x0), generator=gen).to(device) noise = torch.randn(x0.shape, generator=gen).to(device) x_t = (1.0 - t)[:, None] * noise + t[:, None] * x0 target_v = x0 - noise # velocity of the straight noise->data line pred = model.velocity(x_t, t, model.fuse(tokens, img, state)) # Average the velocity MSE over each example's VALID dims, THEN over examples, so a # 2-D PushT frame and a 6-D ALOHA frame weigh the same. Summing instead lets the # higher-DOF embodiment dominate the gradient and starve the other task (measured). return ((((pred - target_v) ** 2) * mask).sum(1) / mask.sum(1)).mean() @torch.no_grad()def sample_action(model: TinyVLA, cond: torch.Tensor, steps: int) -> torch.Tensor: """Sample by integrating the velocity ODE from noise (ch1.5), in standardized space.""" x = torch.randn((cond.shape[0], ACT_DIM), generator=gen).to(device) dt = 1.0 / steps for i in range(steps): t = torch.full((cond.shape[0],), i * dt, device=device) x = x + dt * model.velocity(x, t, cond) # forward Euler along the field return xRead TinyVLA.fuse. We lay the inputs out as one sequence:
[CLS, vision, state, tok_0 … tok_15]. The instruction tokens are embedded; the frozen
image feature and the state are each projected to the model width and become one token
apiece; a learned CLS token leads the sequence. A learned positional embedding is added,
padding tokens are masked out of attention, and a couple of self-attention blocks let every
element read every other. The CLS row that comes out has seen words, pixels, and numbers
at once — that is the fused conditioning vector. velocity is then exactly ch1.5's head:
concatenate the noised action, the sinusoidal flow-time embedding, and the fused vector, and
predict the velocity.
Training on the multi-task pile — and the balancing trap
Training is ch1.5's loop with the flow loss conditioned on fuse(...) instead of a bare
state. But mixing two embodiments hides a trap. A PushT frame has 2 real action dims; an
ALOHA frame has 6. If you sum the masked squared error over the batch and divide by the
number of valid dims — the obvious thing — every ALOHA frame contributes three times the
gradient of a PushT frame, and the 6-DOF task quietly takes over. Measured: with that naive
loss, PushT collapsed to 0.0 while ALOHA learned. The fix is one line — average each
example's error over its valid dims first, then average over examples, so a 2-D frame and
a 6-D frame weigh the same:
# Standardize actions once (masked), then train the fused policy with the flow loss.torch.manual_seed(args.seed) # policy init reproducible, independent of the frozen encoder abovegen.manual_seed(args.seed) # fresh flow-noise stream for trainingx0_all = torch.from_numpy((actions - act_mean) / act_std).to(device) * torch.from_numpy(act_mask_np).to(device)mask_t = torch.from_numpy(act_mask_np).to(device)tokens_t = torch.from_numpy(tokens_np).to(device)feats_t = torch.from_numpy(image_feats).to(device)states_t = torch.from_numpy(states).to(device)policy = TinyVLA(len(vocab), feature_dim, args.model_dim, args.layers, args.heads, args.hidden, STATS).to(device)optimizer = torch.optim.Adam([p for p in policy.parameters() if p.requires_grad], lr=args.lr)shuffle = torch.Generator().manual_seed(args.seed + 1) # torch-side RNG for batch ordertrain_loss, step = float("nan"), 0for epoch in range(args.epochs): epoch_loss, nb = 0.0, 0 for batch in torch.randperm(len(states_t), generator=shuffle).split(args.batch_size): loss = flow_loss(policy, x0_all[batch], mask_t[batch], tokens_t[batch], feats_t[batch], states_t[batch]) optimizer.zero_grad() loss.backward() optimizer.step() epoch_loss, nb = epoch_loss + loss.item(), nb + 1 if args.rerun: rr.set_time("step", sequence=step) rr.log("policy/loss/train", rr.Scalars([loss.item()])) step += 1 train_loss = epoch_loss / nb if epoch % 10 == 0 or epoch == args.epochs - 1: print(f"epoch {epoch:3d} flow_mse {train_loss:.5f}")(That trap is exercise ex3 — a self-contained bug-hunt on the masked loss.)
The frozen encoder is part of the policy
There is an honest subtlety in the eval. To roll out, the policy needs an image feature
for every live frame — so it needs the encoder. Chapter 1.7 saved the features, not the
encoder weights and not the frames. So vla.py rebuilds ch1.7's frozen encoder — same
class, same seed, same first random draw — and it comes out bit-identical (the envs and
experts touch no torch RNG, so the encoder is the first thing to draw after the seed). The
lesson is real: a frozen encoder is not a preprocessing step you can throw away — it is part
of the policy's input contract, and you carry it all the way to deployment.
# The frozen encoder, rebuilt IDENTICALLY to ch1.7's (same class, same seed, same first# torch-RNG draw after set_seed — envs/experts touch no torch RNG, verified). At eval we# render a live frame and must project it exactly as training did; ch1.7 saved the# features but not the weights, so the recipe IS the interface we carry forward.class FrozenVisionEncoder(nn.Module): def __init__(self, width: int, out_dim: int) -> None: super().__init__() self.stem = nn.Sequential( nn.Conv2d(3, width, 3, stride=2, padding=1), nn.ReLU(), # 96 -> 48 nn.Conv2d(width, 2 * width, 3, stride=2, padding=1), nn.ReLU(), # 48 -> 24 nn.Conv2d(2 * width, 4 * width, 3, stride=2, padding=1), nn.ReLU(), # 24 -> 12 nn.AdaptiveAvgPool2d(1), ) self.head = nn.Linear(4 * width, out_dim) for p in self.parameters(): p.requires_grad_(False) # FROZEN — never trained here or in ch1.7 self.eval() @torch.no_grad() def forward(self, images_uint8: torch.Tensor) -> torch.Tensor: x = images_uint8.to(torch.float32).permute(0, 3, 1, 2) / 127.5 - 1.0 return self.head(self.stem(x).flatten(1)) encoder = FrozenVisionEncoder(CONV_WIDTH, feature_dim).to(device) # FIRST torch-RNG use: matches ch1.7 def encode_instruction(text: str) -> np.ndarray: """Word-level tokenizer (ch1.7): [BOS] ids [EOS], OOV -> <unk>, pad/truncate.""" stoi = {w: i for i, w in enumerate(vocab)} ids = [stoi["<bos>"]] + [stoi.get(w, stoi["<unk>"]) for w in text.split()] + [stoi["<eos>"]] ids = ids[:MAX_TOKENS] + [PAD_ID] * (MAX_TOKENS - len(ids)) return np.asarray(ids[:MAX_TOKENS], dtype=np.int64)What it can do, and what it can't
Run it — --seed 0 --device cpu, about 2.2 minutes on a CPU laptop:
# Loss measures velocity fits; rollouts measure the task. At each env step: render a# frame, encode it with the FROZEN encoder, fuse with the (fixed, per-task) instruction# and current state, and SAMPLE an action by ODE integration. Report a Wilson 95%# interval on the success rate — VLA success is noisy and a bare % lies (ch1.6).TASKS = [(PushTEnv, 0, 2), (AlohaCubeEnv, 1, 6)] # (env class, task_id, real action dims) def wilson_ci(k: int, n: int) -> tuple[float, float]: """95% Wilson score interval for k successes in n trials (ch1.6; numpy-free math).""" if n == 0: return (0.0, 1.0) p, z = k / n, Z95 denom = 1.0 + z * z / n center = (p + z * z / (2 * n)) / denom half = (z / denom) * math.sqrt(p * (1.0 - p) / n + z * z / (4 * n * n)) return (max(0.0, center - half), min(1.0, center + half)) @torch.no_grad()def rollout(model: TinyVLA, env_cls, tok_row: np.ndarray, act_dim: int, ep_seed: int) -> tuple[bool, float]: env = env_cls() obs = env.reset(ep_seed) gen.manual_seed(ep_seed) # seed the sampler from the episode: reproducible AND order-independent tok = torch.from_numpy(tok_row).to(device).unsqueeze(0) done, info, ret = False, {}, 0.0 while not done: feat = encoder(torch.from_numpy(env.render_frame(IMG_HW, IMG_HW)[None]).to(device)) cond = model.fuse(tok, feat, torch.from_numpy(obs[None]).to(device)) x = sample_action(model, cond, args.flow_steps) action = (x * act_std_t + act_mean_t)[0, :act_dim].cpu().numpy().clip(-1.0, 1.0) obs, reward, done, info = env.step(action) ret += reward return bool(info["success"]), ret def evaluate(model: TinyVLA, env_cls, task_id: int, act_dim: int, episodes: int, tag: str) -> tuple[int, int, float]: instruction = manifest["tasks"][task_id]["templates"][0] # a fixed, held-in instruction for this task tok_row = encode_instruction(instruction) outcomes = [rollout(model, env_cls, tok_row, act_dim, 10_000 + args.seed + ep) for ep in range(episodes)] k = sum(s for s, _ in outcomes) mean_return = float(np.mean([r for _, r in outcomes])) lo, hi = wilson_ci(k, episodes) # mean_return separates a policy that DRIVES TOWARD the goal (0% success but better # shaped reward) from one that wanders — the honest learning signal when success is 0. print(f"eval[{tag:16s}] {manifest['tasks'][task_id]['name']:6s}: success {k}/{episodes} = " f"{k / episodes:.2f} 95% CI [{lo:.2f}, {hi:.2f}] mean_return {mean_return:.2f}") if args.rerun: rr.log(f"eval/{tag}/success_rate", rr.Scalars([k / episodes])) rr.log(f"eval/{tag}/ci", rr.Scalars([lo, hi])) rr.log(f"eval/{tag}/mean_return", rr.Scalars([mean_return])) return k, episodes, mean_return # A fixed random-init reference (ch1.5 pattern): shows the trained number is signal.torch.manual_seed(args.seed + 2)baseline = TinyVLA(len(vocab), feature_dim, args.model_dim, args.layers, args.heads, args.hidden, STATS).to(device)policy.eval()kp, np_, ret_p = evaluate(policy, PushTEnv, 0, 2, args.eval_episodes, "trained")kb, nb_, ret_b = evaluate(baseline, PushTEnv, 0, 2, args.eval_episodes, "untrained")ka, na, ret_a = evaluate(policy, AlohaCubeEnv, 1, 6, max(1, args.eval_episodes // 2), "trained") # Fused-attention viz: how much the CLS token attends to vision vs state vs each word,# read off the last block for one PushT example — the picture of what the fusion uses.if args.rerun: tok0 = torch.from_numpy(encode_instruction(manifest["tasks"][0]["templates"][0])).to(device).unsqueeze(0) policy.fuse(tok0, feats_t[:1], states_t[:1]) rr.log("fusion/cls_attention", rr.BarChart(policy.blocks[-1].last_attn[0].cpu().double().numpy()))The measured result, seed 0, default config:
- PushT: trained 0.58, Wilson 95% CI [0.32, 0.81], versus an untrained 0.0 (CI [0.00, 0.24]); mean return −48 against −104. Across seeds 0/1/2 the trained rate is 0.58 / 0.58 / 0.42 — above the untrained 0.0 every time. The from-scratch VLA learns PushT.
- ALOHA: 0.0 on every seed. The bimanual handoff needs coordinated multi-step planning (the reason ch1.3's ACT chunked its actions); a tiny shared-capacity policy sampling one action at a time, through a random vision encoder, cannot do it.
So far this looks like a modest win. Now the uncomfortable question a good VLA engineer always asks: is the policy actually using its eyes?
Break It: --break blind
--break blind zeros the image feature at both train and eval. The policy still gets
the instruction and the state — it just never sees anything. Predict the PushT success
before you read on (that is exercise ex1).
Measured: PushT is unchanged — 0.58 blind versus 0.58 sighted, the mean return if
anything slightly better. Zeroing the camera did nothing, because PushT's answer is already
in the state vector (chapter 1.1 solved PushT from state alone), and a random-init vision
feature had nothing to add. The fusion/cls_attention bar in the .rrd tells the same
story: the read-out token barely weights the vision slot.
This is not a failure of the code — the code is a correct, honest VLA. It is a failure of the ingredients. A from-scratch random encoder is a fixed projection of the pixels, not perception. It preserves coarse layout (ch1.7), but nothing about it is aligned to objects or to language, so a policy trained on it learns to ignore it whenever the state suffices — and to fail whenever the state does not (ALOHA).
Why you'd reach for SmolVLA (the Scale Lab)
That gap is the entire argument for adapt-pretrained over from-scratch. A real VLA — SmolVLA, OpenVLA — replaces two of our from-scratch parts with pretrained ones: a vision backbone (SigLIP/DINOv2) whose features are aligned to objects and language, and a subword tokenizer from a pretrained LM. The action head can stay a flow-matching expert — the very mechanism you built in 1.5. The Scale Lab fine-tunes SmolVLA on a consumer GPU and measures the other half of this tradeoff. From-scratch taught you every moving part; adapt-pretrained is what makes the vision and the language actually load-bearing.
A note on what fit
An honest from-scratch VLA — tokenizer reuse, a rebuilt frozen encoder, a multi-head attention backbone, a conditioned flow head, masked multi-task training, and a two-task error-bar eval — fits in one 449-line file (hard cap 450, target 400). It is tight. Two things made it fit: there is no 2-D toy here (unlike ch1.5), and there is no ONNX export (a full VLA does not fit the stateless demo contract anyway — the browser panel is blocked on the same flow-sampler contract v2 as the ch1.5 policy). If a future teaching-pass needs more room, the honest cut is the ALOHA eval (keep multi-task training, evaluate only PushT) — not a file split.
What we cut
- Pretrained everything. Real VLAs use a pretrained vision backbone and a pretrained-LM subword tokenizer. Ours are frozen-random and 46-words-fixed (ch1.7). That is the Scale Lab, and the reason this policy ignores its vision.
- Action chunking. We sample a single action per step (ch1.5). ACT (ch1.3) and real VLAs predict a chunk; that is likely the missing ingredient for ALOHA.
- A bigger backbone, more data, more tasks. Scale knobs are visible flags (
--model_dim,--layers,--heads,--epochs,--episodes_per_task).
Read the real thing
The Scale Lab is not hypothetical — it is a real policy you can read line for line.
huggingface/lerobot, pinned here at v0.4.4 (commit 8fff0fd), ships SmolVLA
under src/lerobot/policies/smolvla/. It is the exact same three parts you just built —
a fusion backbone, a flow head, a tokenizer — with our from-scratch pieces swapped for
pretrained ones. Read it in three passes, against your vla.py.
The fusion backbone. Your TinyVLA.fuse (the model region) lays out
[CLS, vision, state, tok_0..tok_15] and runs a couple of from-scratch attention Blocks;
the vision token is vision_proj of a frozen random CNN (the vision_language
region's FrozenVisionEncoder). SmolVLA's VLAFlowMatching.embed_prefix() in
src/lerobot/policies/smolvla/modeling_smolvla.py does the identical lay-out — image
tokens, language tokens, a state_proj — but the image tokens come from a pretrained
SmolVLM2 backbone (SmolVLMWithExpertModel in
src/lerobot/policies/smolvla/smolvlm_with_expert.py loads
HuggingFaceTB/SmolVLM2-500M-Video-Instruct via AutoModelForImageTextToText; the vision
tower is its SigLIP vision_model + connector). That is the one change that makes vision
load-bearing: features already aligned to objects and language, where ours were a fixed
random projection (--break blind).
The action expert. Your flow_loss regresses the straight-line velocity
target_v = x0 - noise, and sample_action integrates it with forward Euler in
flow_steps steps — the ch1.5 mechanism. SmolVLA's VLAFlowMatching builds
x_t = t*noise + (1-t)*actions, regresses u_t = noise - actions with an MSE on the
predicted velocity, and its sample_actions() / denoise_step() integrate noise→action
over num_steps (default 10, dt = -1/num_steps) — your loop, on a bigger field. What
they add: the velocity head is a half-width transformer "expert" that cross-attends to
the frozen VLM prefix (forward_cross_attn_layer, attention_mode) and predicts a
chunk of 50 actions, not the single action our sampler draws.
The tokenizer. Your encode_instruction (the vision_language region) maps words
through ch1.7's fixed 46-word vocab. SmolVLA's processor_smolvla.py runs a
TokenizerProcessorStep(tokenizer_name=config.vlm_model_name, ...) — the pretrained
subword tokenizer shipped with SmolVLM2 (tokenizer_max_length=48), plus a
SmolVLANewLineProcessor quirk-fix. Any instruction tokenizes; no OOV cliff at 46 words.
Read these, in order. modeling_smolvla.py's embed_prefix and sample_actions
first — your fuse and sample_action, grown up. Then smolvlm_with_expert.py, to see
the pretrained backbone our FrozenVisionEncoder stands in for. That backbone is the
whole reason --break blind would not be a no-op for SmolVLA — and the whole reason to
adapt-pretrained when performance matters.
Exercises
- ex1 (predict-then-run): does
--break blindhurt PushT? (Measure that vision is not load-bearing.) - ex2 (predict-then-run): which task does the tiny VLA learn — PushT or ALOHA?
- ex3 (bug-hunt): the multi-task masked loss — fix the sum-weighting that lets the 6-DOF embodiment dominate.
- ex4 (code-completion): write the scaled dot-product attention at the heart of the fusion backbone.