The escape hatch, removed
Way back in chapter 1.1 you cloned PushT from a ten-number state vector (pusher xy,
block pose, target pose) and it just worked. A three-layer MLP, plain MSE, done. Then in
chapter 1.8 you built a vision-language-action policy, gave it a camera, and ran
--break blind: zeroing the image feature changed PushT success by nothing. The policy
never used its eyes. We told you why: PushT is solvable from the state, and a random-init
encoder had nothing to add, but that was an argument, not a measurement. This chapter turns
it into one.
Here is the whole move: take away the state. The policy gets one input, a live 64×64 top-down frame, and must produce the pusher velocity from pixels alone. No coordinates, no cheating. The same BC loop from chapter 1.1 (load pairs, fit a network, roll it out, export to ONNX), but the observation is now an image. And the moment the state is gone, one thing you could ignore in chapter 1.8 becomes the entire chapter: the quality of the encoder that turns pixels into features.
pixels.py runs that BC loop twice, changing exactly one part, the encoder:
- ALIGNED: a tiny ViT that has been contrastively pre-aligned to the scene geometry (the chapter 5.2 recipe, re-contained and compact). Frozen; only a thin adapter and the policy head train.
- RANDOM: the identical ViT architecture, never aligned. A fixed random projection of the pixels, exactly the kind of encoder chapter 1.8's VLA looked through and ignored.
Everything downstream (the adapter, the head, the BC loss, the rollout) is held fixed. The only variable is what the encoder learned. That is the experiment.
The misconception this kills
"Pixels → action is just behavior cloning with more inputs."
It is not, and the reason is the whole lesson. In chapter 1.1 the input was a state vector that already was the answer: block position, target position, a straight line between them. Any encoder, even the identity, hands the MLP everything it needs. Swap in pixels and that free lunch disappears. Raw pixels do not linearly encode "where is the block relative to the target"; something has to extract that. From pixels with no state, the encoder's quality is not a detail of the input: it is the entire policy. A random projection scrambles the geometry; the BC head, no matter how well trained, is cloning from noise.
The tiny ViT (the shared backbone, re-contained)
Every chapter re-contains its own backbone. The repetition is the lesson. The encoder is the
same block shape chapter 1.8's VLA used: a stride-8 conv patch-embed turns the 64×64 frame
into an 8×8 grid of 64 patch tokens, a learned CLS token leads the sequence, learned
positions are added, and a few pre-norm self-attention blocks let every patch talk to every
other. The CLS row that comes out is the image feature the policy conditions on.
# A tiny ViT — the SAME block shape ch1.8's VLA uses, re-derived here (chapters re-contain their# backbone; the repetition is the lesson). Stride-8 conv -> 64 tokens, a learned CLS + positions,# a few pre-norm self-attention blocks, read the CLS.class Block(nn.Module): """One pre-norm transformer block: multi-head self-attention + an MLP, 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 attention over patches, for the saliency viz def forward(self, x: 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) attn = ((q @ k.transpose(-2, -1)) / math.sqrt(hd)).softmax(dim=-1) self.last_attn = attn[:, :, 0, 1:].mean(1).detach() # (B, 64): CLS over the patch grid x = x + self.proj((attn @ v).transpose(1, 2).reshape(B, L, dim)) return x + self.mlp(self.ln2(x)) class ViTEncoder(nn.Module): """(B,64,64,3) float pixels -> (B, dim) CLS feature (random until aligned). Subtracts the background (frame_mean) FIRST so the ViT sees the block, not the table.""" def __init__(self, dim: int, depth: int, heads: int, frame_mean: torch.Tensor) -> None: super().__init__() self.patch = nn.Conv2d(3, dim, kernel_size=PATCH, stride=PATCH) # -> (B, dim, 8, 8) self.cls = nn.Parameter(torch.zeros(1, 1, dim)) self.pos = nn.Parameter(0.02 * torch.randn(1, 1 + (IMG_HW // PATCH) ** 2, dim)) self.blocks = nn.ModuleList([Block(dim, heads) for _ in range(depth)]) self.norm = nn.LayerNorm(dim) self.register_buffer("frame_mean", frame_mean) # background, subtracted below def forward(self, images: torch.Tensor) -> torch.Tensor: x = images / 127.5 - 1.0 # (B, 64, 64, 3) in [-1, 1] x = (x - self.frame_mean).permute(0, 3, 1, 2) # center on background, then (B, 3, 64, 64) x = self.patch(x).flatten(2).transpose(1, 2) # (B, 64, dim) patch tokens x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], dim=1) + self.pos for blk in self.blocks: x = blk(x) return self.norm(x[:, 0]) # the CLS row: the image feature class PixelPolicy(nn.Module): """The deployable policy: a FROZEN ViT encoder + a thin trained adapter + head. forward takes the flat image as the contract-v1 observation, so the whole pixels->action path exports as one ONNX graph (the encoder is INSIDE). Normalization rides in buffers.""" def __init__(self, encoder: ViTEncoder, dim: int, hidden: int, feat_mean, feat_std, act_min, act_range) -> None: super().__init__() self.encoder = encoder self.adapter = nn.Linear(dim, hidden) # the thin adapter onto the frozen features self.head = nn.Linear(hidden, ACT_DIM) for name, stat in [("feat_mean", feat_mean), ("feat_std", feat_std), ("act_min", act_min), ("act_range", act_range)]: self.register_buffer(name, torch.as_tensor(stat, dtype=torch.float32)) def head_forward(self, feat_std: torch.Tensor) -> torch.Tensor: # standardized feature -> (B, hidden) -> tanh action in [-1, 1] -> raw velocity a = torch.tanh(self.head(torch.relu(self.adapter(feat_std)))) return (a + 1.0) / 2.0 * self.act_range + self.act_min def forward(self, obs: torch.Tensor) -> torch.Tensor: # (B, 12288) flat image -> (B, 2) action. The frozen encoder is INSIDE the policy. feat = self.encoder(obs.reshape(-1, IMG_HW, IMG_HW, 3)) return self.head_forward((feat - self.feat_mean) / self.feat_std)Read PixelPolicy. Its forward takes the flat image as the observation (the encoder
lives inside the policy), so the entire pixels→action path exports as one ONNX graph under
tensor contract v1 (the encoder reshapes [1, 12288] back into a frame). That is the honest
consequence of chapter 1.8's lesson that a frozen encoder is part of the policy's input
contract: here we carry it literally into the deployed graph.
Aligning the encoder (the ch5.2 recipe, compact)
Chapter 5.2 taught contrastive alignment in full. We re-contain a compact version, because this chapter needs the product (an aligned encoder) more than a second full treatment. The objective is a symmetric InfoNCE: for each frame in a batch, pull its image feature onto its own scene geometry and push it off every other example's. The encoder is never told the action, only that these pixels and this geometry belong together, and that is enough to make its features carry where the block and pusher are.
# THE ch5.2 RECIPE, re-contained and COMPACT (5.2 teaches it in full). Symmetric InfoNCE pulls each# image's CLS feature onto its own scene state and off every other example's. The encoder is told# only "these pixels and this geometry belong together" — enough to make the features carry WHERE# the block and pusher are. Deterministic from the seed (CPU), so re-running reproduces it — no git.def align_encoder(encoder: ViTEncoder, epochs: int) -> float: # the contrastive PARTNER tower: projects the 10-D state into the ViT's feature space, pulling # image features onto the scene geometry. Thrown away after alignment (only the encoder is kept). state_enc = nn.Sequential(nn.Linear(PushTEnv.OBS_DIM, 128), nn.ReLU(), nn.Linear(128, args.dim)).to(device) params = list(encoder.parameters()) + list(state_enc.parameters()) opt = torch.optim.Adam(params, lr=args.lr) shuffle = torch.Generator().manual_seed(args.seed + 7) loss_val = float("nan") for epoch in range(epochs): for batch in torch.randperm(num_frames, generator=shuffle).split(args.batch_size): if len(batch) < 2: # InfoNCE needs negatives — a singleton batch has none continue img = nn.functional.normalize(encoder(frames_t[batch]), dim=1) st = nn.functional.normalize(state_enc(states_t[batch]), dim=1) logits = img @ st.T / TEMP # (b, b) image-state similarities labels = torch.arange(len(batch), device=device) # the diagonal is the match loss = 0.5 * (nn.functional.cross_entropy(logits, labels) + nn.functional.cross_entropy(logits.T, labels)) opt.zero_grad() loss.backward() opt.step() loss_val = loss.item() if args.rerun: rr.log("align/infonce", rr.Scalars([loss_val])) return loss_valA note on faithfulness to 5.2: there the partner is a text tower; here it is the state, because on single-instruction PushT the geometry, not the words, is the signal that makes pixel features control-relevant. The mechanism is identical (a tiny ViT, a partner tower, symmetric InfoNCE); only the partner differs. And, exactly like chapter 1.8's frozen encoder, the aligned encoder is rebuilt deterministically from the seed: CPU alignment is reproducible, so re-running it reproduces the weights with no checkpoint binary in git.
Cloning from frozen features
With the encoder aligned and frozen, the rest is chapter 1.1 with the observation swapped. We featurize every frame once, standardize, and fit the adapter + head with plain MSE. The random-encoder run is byte-for-byte the same loop over a ViT that skipped alignment.
# ch1.1's BC loop, obs swapped for frozen features. Freeze the encoder, featurize every frame ONCE# (fast: the head trains on vectors, not pixels), standardize, fit the adapter+head with plain MSE.# `--train_encoder` is THE TRAP: it unfreezes the encoder end-to-end, where the tiny set overfits.def build_and_train(encoder: ViTEncoder, trainable: bool, tag: str) -> tuple[PixelPolicy, float]: feats = encode_all(encoder) # (N, dim) frozen features feat_mean, feat_std = feats.mean(0), feats.std(0).clamp_min(1e-4) policy = PixelPolicy(encoder, args.dim, args.hidden, feat_mean.cpu().numpy(), feat_std.cpu().numpy(), act_min, act_range).to(device) if trainable: # TRAP: encoder joins the optimizer and overfits the pixels params, feats_std = list(policy.parameters()), None for p in policy.encoder.parameters(): p.requires_grad_(True) else: # frozen: only the adapter+head learn, on precomputed features for p in policy.encoder.parameters(): p.requires_grad_(False) params = list(policy.adapter.parameters()) + list(policy.head.parameters()) feats_std = (feats - feat_mean) / feat_std opt = torch.optim.Adam(params, lr=args.lr) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.bc_epochs) shuffle = torch.Generator().manual_seed(args.seed + 1) loss_val = float("nan") for epoch in range(args.bc_epochs): policy.train() for batch in torch.randperm(num_frames, generator=shuffle).split(args.batch_size): if trainable: # recompute features through the live encoder pred = policy(frames_t[batch].reshape(len(batch), -1)) else: # head-only on the standardized frozen features pred = policy.head_forward(feats_std[batch]) loss = nn.functional.mse_loss(pred, actions_t[batch]) opt.zero_grad() loss.backward() opt.step() loss_val = loss.item() sched.step() if args.rerun: rr.log(f"bc/{tag}/loss", rr.Scalars([loss_val])) return policy, loss_valWhat the measurement says
Run it: --seed 0 --device cpu, wall-clock ~2.95 min (measured on cpu-laptop; see
curriculum/common/wallclock.csv, and the banner prints it at startup). Both policies
roll out from pixels alone, and we report a Wilson 95% interval (chapter 1.6), because
pixel-BC success is noisy and a bare percentage lies:
# Rollouts from PIXELS. Each step renders the live 64x64 frame, flattens it to the contract-v1# observation, and steps the env with the policy's action — the state MuJoCo tracks is never shown# to the policy. Report a Wilson 95% interval (ch1.6): pixel-BC success is noisy and a bare % lies.def wilson_ci(k: int, n: int) -> tuple[float, float]: 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(policy: PixelPolicy, ep_seed: int, record: bool = False): policy.eval() pdev = next(policy.parameters()).device # export_onnx moves the policy to cpu in place — follow it env = PushTEnv() env.reset(ep_seed) done, info, ret, traj = False, {}, 0.0, [] while not done: frame = env.render_frame(IMG_HW, IMG_HW) obs = torch.from_numpy(frame).to(pdev).float().reshape(1, -1) # (1, 12288) action = policy(obs)[0].cpu().numpy() if record: px, py = env.pusher_pos tx, ty, tyaw = env.tee_pose traj.append([round(float(px), 4), round(float(py), 4), round(float(tx), 4), round(float(ty), 4), round(float(tyaw), 4)]) _, reward, done, info = env.step(action) ret += reward return bool(info["success"]), ret, traj def evaluate(policy: PixelPolicy, tag: str) -> dict: outcomes = [rollout(policy, 10_000 + args.seed + ep) for ep in range(args.eval_episodes)] k = sum(s for s, _, _ in outcomes) lo, hi = wilson_ci(k, args.eval_episodes) mean_return = float(np.mean([r for _, r, _ in outcomes])) print(f"eval[{tag:8s}] pixels-only: success {k}/{args.eval_episodes} = {k / args.eval_episodes:.2f} " f"95% CI [{lo:.2f}, {hi:.2f}] mean_return {mean_return:.2f}") if args.rerun: rr.log(f"eval/{tag}/success_rate", rr.Scalars([k / args.eval_episodes])) rr.log(f"eval/{tag}/ci", rr.Scalars([lo, hi])) return {"success_rate": k / args.eval_episodes, "ci_lo": lo, "ci_hi": hi, "mean_return": mean_return}The thing to watch is the direction, not an absolute number. The reproducible one is a
probe on the frozen features: fit a tiny linear map from the encoder's features to the
expert's action, and compare the held-out error for the aligned encoder against a random one
(probe_val_mse, aligned vs random). When the alignment has made the features carry the geometry,
that action-probe error is lower. Aligned beats random on every seed: +0.028 / +0.040 /
+0.054. That is the gated headline, and it stands on its own: aligned features are measurably
more control-useful than random ones.
What floors, and why (said once, then trusted for the rest of the chapter). The closed-loop
success_rate is a harder bar than the probe, and at free-tier scale it floors at 0/12 for
both encoders: the rollout gap is 0.0. That is not the experiment failing. It is the ceiling of
a tiny from-scratch encoder aligned to a few hundred frames, nowhere near SigLIP quality, cloning
a non-Markovian expert one action at a time from a single 64×64 frame. Two things follow, and we
will not repeat them each section. First, the reproducible signal is the probe, not a rollout
win, so that is what the chapter gates (chapter 1.6). Second, absolute pixel-BC success is small
and platform-sensitive: MuJoCo rasterization is not bitwise across CPU architectures (chapter
1.8), so we report a direction under a Wilson interval, never a promised number. Driving PushT
end to end from a scaled backbone is the Scale Lab. That honest ceiling is the lesson: a
from-scratch encoder is a fixed, weak projection of the pixels, which is the whole reason the real
answer (the read-the-real-thing) is a backbone aligned at internet scale.
The payoff to --break blind
This is the measured answer to chapter 1.8's cliffhanger. There, zeroing vision was a no-op, because the state made vision redundant and the random encoder added nothing. Here the state is gone, and the two encoders sit at opposite ends of the same experiment: the aligned features carry enough of the geometry that a linear probe reads the expert's action out of them better than from random features, on every seed. That gap is the thing chapter 1.8 could only assert: that a pretrained, aligned backbone is what makes vision actually matter. When the state cannot bail you out, the encoder is the whole policy, and the quality of its features is the whole ballgame.
Break it yourself: the trainable-encoder trap
You wired the encoder→policy path. The obvious next thought: why freeze the encoder? Won't
training the whole thing end-to-end beat training a thin adapter? Predict the answer before
you run --train_encoder (that is ex2), then run it. At free-tier scale it is a trap: a
from-scratch ViT has enough capacity to memorize the handful of frames it trained on, so the
pixels→action map overfits: a tiny training loss sitting over a worse held-out fit (the tell is
that overfit gap, not a rollout number). Then explain it to yourself: chapter 1.1 trained
its whole network end-to-end and was fine. Why is this different? (A ten-number state cannot
be memorized into a lookup table; a tiny frame set can. The alignment already did the
expensive, transferable work. Leave it frozen.)
Read the real thing
The definitive version of "a pretrained, aligned vision backbone feeds the action decoder"
is OpenVLA (openvla/openvla; the author pins the verified commit). Where our policy runs
one frozen CLS feature through a thin nn.Linear adapter into a BC head, OpenVLA runs images
through Prismatic's fused SigLIP + DINOv2 backbone (vision features aligned to objects and
language at internet scale), projects them into an LLM's token space, and decodes actions. Read
it for the one structural echo of this chapter: the backbone is frozen (or only lightly
adapted), because the alignment is the expensive, transferable part: the same reason our
--train_encoder trap loses to the frozen aligned encoder, and the same aligned-vs-random gap
this chapter measures, scaled up a million-fold. (Alternative: Physical-Intelligence/openpi's
pi0, a PaliGemma/SigLIP image path into a flow-matching action expert.)
What we cut
- A real aligned backbone. Ours is a tiny from-scratch ViT aligned to sim state on a few hundred frames: modestly better than random, not SigLIP. The whole point of OpenVLA / pi0 is that the alignment is pretrained at scale. That is the read-the-real-thing.
- Action chunking and history. We clone a single action from a single frame (chapter 1.1's shape). ACT (chapter 1.3) and real visuomotor policies predict a chunk from a short history, likely what a harder task would need.
- Language. PushT has one instruction, so alignment here is image↔state. Chapter 5.2's image↔text alignment is what you would reach for on a multi-instruction task.
Exercises
- ex1 (predict-then-run): aligned vs random encoder: whose frozen features are more control-useful? (The gated signal is the action-probe direction; the encoder is load-bearing once the state is gone.)
- ex2 (predict-then-run + the trap): frozen aligned encoder vs
--train_encoder. Predict, run, and self-explain why unfreezing overfits here but not in chapter 1.1. - ex3 (code-completion): write the symmetric InfoNCE at the heart of the alignment.