A VLA is trained like any imitation policy: collect demonstrations of a task described in words, then fit the network so that its action matches the demonstrator’s on every recorded frame. The MiniVLA of Lesson 22.2 is trained this way on 3000 demonstrations of the tabletop, in about 14 minutes on 16 CPU cores, and evaluated by running it in MuJoCo on layouts it has never seen.
The first version failed completely: 1 success in 100 episodes, for all three vision encoders, after reaching a training loss lower than that of models that later succeeded 96% of the time. This lesson reproduces that failure and finds its cause, a shortcut through the pusher’s velocity that the network took because the data allowed it. The cause is a standard one in imitation learning, and finding it took the same steps a research project would: an oracle that rules out vision, an offline-against-online comparison, and an intervention on one input.
One held-out task, “get the green puck into the right zone”, run for 6 s from the same start by the demonstrator and by two patch MiniVLAs that differ only in whether the pusher’s velocity is part of their input. The demonstrator (left) and the position-only model (right) go around the green puck and push it into the zone. The model that sees its velocity (centre) draws a smooth curve of about the right shape and length, past the puck, without touching it: it reproduces how the demonstrator moves, not where it moves to.
Each demonstration samples one of the 216 training instructions (9 object-zone pairs × 6 templates × 4 name variants; Lesson 22.5 holds out the rest), a layout, and runs the demonstrator of mjcourse.tabletop.expert for up to 12 s, recording at 10 Hz the 64 × 64 top-camera image $o_t$, the instruction $\ell$, the pusher’s state $s_t$ and the demonstrator’s action $a^*_t \in [-1,1]^2$ (velocity as a fraction of 0.2 m/s). Of 3000 episodes, 2995 succeed and are kept: 169 337 frames. Behaviour cloning minimizes
with AdamW (learning rate $3\times10^{-4}$, weight decay $10^{-4}$, cosine decay), batches of 64, gradient clipping at 1, and random image shifts of up to 4 pixels, for 30 000 steps (11 passes over the frames).
Output (1) of the diagnosis script collects every trained variant. Success is the share of episodes in which the named puck ends within 5 cm of its zone’s centre within 12 s, on layouts from seeds the demonstrations never used (100 for training instructions, 60 for each held-out language split of Lesson 22.5).
| Vision | State input | Seed | Final loss | Success, training instructions |
|---|---|---|---|---|
| patch | position and velocity | 0 | 0.0173 | 0.01 |
| keypoint | position and velocity | 0 | 0.0179 | 0.01 |
| oracle (true puck positions) | position and velocity | 0 | 0.0160 | 0.01 |
| patch | position | 0, 1, 2 | 0.0155, 0.0151, 0.0142 | 0.95, 0.92, 0.99 |
| keypoint | position | 0, 1, 2 | 0.0315, 0.0111, 0.0311 | 0.58, 0.96, 0.62 |
| oracle | position | 0, 1, 2 | 0.0068, 0.0047, 0.0067 | 0.99, 1.00, 1.00 |
Three findings, in the order they were obtained.
The velocity models fail with every encoder, including the oracle, which never sees an image. A failure that survives perfect perception is not a perception failure; this one test removed vision from the list of suspects.
The training loss does not order the models by success. The failed velocity models reach 0.016 to 0.018, lower than two keypoint runs (0.031) that succeed about 60% of the time.
Removing the velocity fixes all three encoders. The patch model, chosen for the rest of Level 22, reaches 0.92 to 0.99 over three seeds. The keypoint encoder is unstable: two seeds stall at a loss of 0.031 and one reaches 0.011; a single-seed result would have reported either 0.58 or 0.96.
The pusher is driven by velocity actuators that track their command within a few milliseconds, and the demonstrator’s commands change smoothly. The pusher’s velocity at frame $t$ is therefore almost the demonstrator’s action at frame $t-1$, and that action is almost the action at frame $t$: copying the previous action has an error of 0.071 (output 2), against actions of mean norm 0.58. A network given the velocity can fit the data to that accuracy without learning where any puck is. De Haan et al. (arXiv:1905.11979) call the general phenomenon causal confusion, in which “access to more information can yield worse performance”; Wen et al. (arXiv:2010.14876) describe this instance, an imitator that “learns to cheat by predicting the expert’s previous action”.
Output (2) confirms it offline, on 40 held-out demonstrations:
| Model | Error at each episode’s first frame (pusher at rest) | Error at later frames |
|---|---|---|
| patch, position and velocity | 0.653 | 0.058 |
| patch, position only | 0.079 | 0.090 |
| copying the previous action | (none) | 0.071 |
On later frames the velocity model is the most accurate of the three, better even than copying. At the first frame, where the velocity is zero and carries no information, its error is larger than the actions themselves. Output (3) intervenes: setting the velocity input to zero changes its action by 0.75 on average, and its action lies closer to the velocity it is given (0.072) than the demonstrator’s action does (0.092).
In closed loop, the velocity it is given is its own previous action. Output (4) runs both models from rest on 20 held-out tasks and asks the demonstrator what it would have done at each state the model visits:
| Model | Disagreement with the demonstrator, 0 to 1 s | 1 to 2 s | 2 to 3 s |
|---|---|---|---|
| patch, position and velocity | 0.739 | 1.109 | 1.300 |
| patch, position only | 0.129 | 0.163 | 0.184 |
An offline error of 0.058 becomes 0.74 in the first second and keeps growing: the model starts in the wrong direction, its velocity input then tells it to continue, and it does.
[!recommendation] Before training any imitation policy List every input that is a consequence of the previous action (velocities, previous actions, filtered signals) and ask whether it predicts the next action without the task-relevant inputs. If it does, remove it, or make it less predictive (Wen et al. remove the previous-action information from the representation adversarially). Then measure the first-frame error and the on-policy disagreement, never only the held-out loss on demonstration frames.
INPUT: 3000 demonstrations (mjcourse.vla.data.cached(0, 3000)); a vision encoder; a seed
PROCESS: behaviour cloning for 30 000 steps; closed-loop evaluation on held-out layouts for four language splits
OUTPUT: runs/vla/minivla_VISION_sSEED_pos.pt and .json
```python file=examples/l22_4_train_minivla.py import argparse import json import multiprocessing import os import time from concurrent.futures import ProcessPoolExecutor 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 VLAConfig
FAST = os.environ.get(“MJC_FAST”) == “1” OUT = Path(file).resolve().parents[1] / “runs” / “vla” EPISODES = {“train”: 2, “paraphrase”: 1, “synonym”: 1, “composition”: 1} if FAST else {“train”: 100, “paraphrase”: 60, “synonym”: 60, “composition”: 60}
def evaluate_one(args) -> dict: model_path, split, index = args torch.set_num_threads(1) model = torch.load(model_path, weights_only=False) rng = np.random.default_rng([2026, index]) # layouts the demonstrations never used instructions = language.instructions(split) text, task = instructions[rng.integers(len(instructions))] layout = tt.sample_layout(rng, task) result = train.run_episode(train.VLAPolicy(model, language.default_tokenizer()), text, task, layout) return {“split”: split, “index”: index, “text”: text, “task”: list(task), “success”: bool(result[“success”]), “steps”: result[“steps”]}
def main() -> None: parser = argparse.ArgumentParser() parser.add_argument(“–vision”, default=”keypoint”, choices=[“patch”, “keypoint”, “oracle”]) parser.add_argument(“–seed”, type=int, default=0) parser.add_argument(“–steps”, type=int, default=100 if FAST else 30000) parser.add_argument(“–demos”, type=int, default=20 if FAST else 3000) parser.add_argument(“–state”, type=int, default=2, choices=[2, 4], help=”2: pusher position (the course default); 4: position and velocity (Lesson 22.4, the copycat)”) args = parser.parse_args() OUT.mkdir(parents=True, exist_ok=True) cpus = int(os.environ.get(“SLURM_CPUS_PER_TASK”, min(16, os.cpu_count() or 1))) torch.set_num_threads(cpus) name = (f”minivla_{args.vision}_s{args.seed}” + (“” if args.demos == 3000 else f”_d{args.demos}”) + (“” if args.state == 4 else “_pos”)) episodes = data.cached(0, 3000, “train”)[:args.demos] tokenizer = language.default_tokenizer() arrays = data.to_arrays(episodes, tokenizer, 1) print(f”{name}: {sum(e[‘success’] for e in episodes)} successful demonstrations, {len(arrays[‘images’])} frames”, flush=True) start = time.perf_counter() model, losses = train.fit(arrays, len(tokenizer.vocab), steps=args.steps, seed=args.seed, cfg=VLAConfig(vocab=len(tokenizer.vocab), chunk=1, vision=args.vision, state_dim=args.state), log=lambda m: print(m, flush=True)) seconds = time.perf_counter() - start path = OUT / f”{name}.pt” torch.save(model, path) jobs = [(str(path), split, i) for split, n in EPISODES.items() for i in range(n)] with ProcessPoolExecutor(cpus, mp_context=multiprocessing.get_context(“spawn”)) as pool: episodes_out = list(pool.map(evaluate_one, jobs, chunksize=2)) summary = {split: float(np.mean([e[“success”] for e in episodes_out if e[“split”] == split])) for split in EPISODES} record = {“vision”: args.vision, “seed”: args.seed, “steps”: args.steps, “demos”: args.demos, “state”: args.state, “train_seconds”: seconds, “final_loss”: float(np.mean(losses[-1000:])), “summary”: summary, “episodes”: episodes_out} (OUT / f”{name}.json”).write_text(json.dumps(record, indent=1)) print(json.dumps({k: v for k, v in record.items() if k != “episodes”}), flush=True)
if name == “main”: main()
Demonstrations, the training loop and the closed-loop policy:
```python file=src/mjcourse/vla/data.py
from __future__ import annotations
import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor
import numpy as np
from mjcourse import tabletop as tt
from mjcourse.vla import language
def episode(seed: int, index: int, split: str = "train", variation: tt.Variation | None = None, camera: str = "top"):
rng = np.random.default_rng([seed, index])
text, task = language.instructions(split)[rng.integers(len(language.instructions(split)))]
scene = tt.Tabletop(variation)
scene.reset(tt.sample_layout(rng, task))
images, states, actions, pucks = [], [], [], []
for _ in range(tt.EPISODE_STEPS):
s = scene.state()
a = tt.expert(s, task)
images.append(scene.render(camera, 64))
states.append(s[:4])
pucks.append(np.r_[s[4:6], s[8:10], s[12:14]])
actions.append(a)
scene.step(a)
if scene.success(task) and np.linalg.norm(scene.state()[4 + 4 * tt.OBJECTS.index(task[0]) + 2:][:2]) < 0.02:
break
ok = scene.success(task)
scene.close()
return {"text": text, "task": task, "images": np.array(images), "states": np.array(states, np.float32),
"actions": np.array(actions, np.float32), "pucks": np.array(pucks, np.float32), "success": ok}
def generate(seed: int, n: int, split: str = "train", workers: int | None = None) -> list[dict]:
workers = workers or min(32, os.cpu_count() or 1)
with ProcessPoolExecutor(workers, mp_context=multiprocessing.get_context("spawn")) as pool:
return list(pool.map(episode, [seed] * n, range(n), [split] * n, chunksize=4))
def cached(seed: int, n: int, split: str = "train") -> list[dict]:
import pickle
from pathlib import Path
path = Path(__file__).resolve().parents[3] / "runs" / "vla" / f"demos_{split}_{seed}_{n}.pkl"
if not path.exists():
path.parent.mkdir(parents=True, exist_ok=True)
path.write_bytes(pickle.dumps(generate(seed, n, split)))
return pickle.loads(path.read_bytes())
def to_arrays(episodes: list[dict], tokenizer, chunk: int) -> dict:
images, words, states, targets, masks, pucks, tasks = [], [], [], [], [], [], []
for ep in episodes:
if not ep["success"]:
continue
t_len = len(ep["actions"])
tokens = tokenizer(ep["text"])
padded = np.concatenate([ep["actions"], np.zeros((chunk, 2), np.float32)])
task = [tt.OBJECTS.index(ep["task"][0]), list(tt.ZONES).index(ep["task"][1])]
for t in range(t_len):
images.append(ep["images"][t])
words.append(tokens)
states.append(ep["states"][t])
targets.append(padded[t:t + chunk])
masks.append((np.arange(chunk) + t < t_len).astype(np.float32))
pucks.append(ep["pucks"][t])
tasks.append(task)
return {"images": np.array(images), "words": np.array(words), "states": np.array(states),
"targets": np.array(targets), "masks": np.array(masks), "pucks": np.array(pucks), "tasks": np.array(tasks)}
```python file=src/mjcourse/vla/train.py from future import annotations
import time
import numpy as np import torch from torch.nn import functional as F
from mjcourse import tabletop as tt from mjcourse.vla.model import MiniVLA, VLAConfig
STATE_SCALE = np.array([1 / 0.3, 1 / 0.3, 1 / 0.2, 1 / 0.2], np.float32) # pusher x, y (m), vx, vy (m/s) to about +-1
def _batch(arrays, idx, shift: int, state_dim: int = 4): img = torch.as_tensor(arrays[“images”][idx]).permute(0, 3, 1, 2).float() / 255.0 if shift: img = F.pad(img, (shift,) * 4, mode=”replicate”) dx, dy = np.random.randint(0, 2 * shift + 1, 2) img = img[:, :, dy:dy + 64, dx:dx + 64] states = (arrays[“states”][idx] * STATE_SCALE)[:, :state_dim] return (img, torch.as_tensor(arrays[“words”][idx]), torch.as_tensor(states), torch.as_tensor(arrays[“targets”][idx]), torch.as_tensor(arrays[“masks”][idx]), torch.as_tensor(arrays[“pucks”][idx]))
def fit(arrays: dict, vocab: int, steps: int = 6000, batch: int = 64, lr: float = 3e-4, seed: int = 0, shift: int = 4, cfg: VLAConfig | None = None, log=None) -> tuple[MiniVLA, list[float]]: torch.manual_seed(seed) np.random.seed(seed) model = MiniVLA(cfg or VLAConfig(vocab=vocab)) opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, steps) n, losses, start = len(arrays[“images”]), [], time.perf_counter() for step in range(steps): idx = np.random.randint(0, n, batch) img, words, state, target, mask, pucks = batch(arrays, idx, shift, model.cfg.state_dim) pred = model(img, words, state, pucks) loss = (((pred - target) ** 2).sum(-1) * mask).sum() / mask.sum() opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm(model.parameters(), 1.0) opt.step() sched.step() losses.append(float(loss.detach())) if log and (step + 1) % 1000 == 0: log(f”step {step + 1}: loss {np.mean(losses[-1000:]):.4f} ({time.perf_counter() - start:.0f} s)”) return model, losses
class VLAPolicy: def init(self, model: MiniVLA, tokenizer, ensemble: bool = False, decay: float = 0.5): self.model, self.tokenizer, self.ensemble, self.decay = model.eval(), tokenizer, ensemble, decay self.history: list[tuple[int, np.ndarray]] = []
def reset(self) -> None:
self.history = []
@torch.no_grad()
def chunk(self, image: np.ndarray, text: str, state: np.ndarray) -> np.ndarray:
img = torch.as_tensor(image).permute(2, 0, 1).float().unsqueeze(0) / 255.0
words = torch.as_tensor(self.tokenizer(text)).unsqueeze(0)
s = torch.as_tensor((state[:4] * STATE_SCALE)[:self.model.cfg.state_dim].astype(np.float32)).unsqueeze(0)
pucks = None
if self.model.cfg.vision == "oracle": # privileged: diagnosis only
pucks = torch.as_tensor(state[[4, 5, 8, 9, 12, 13]].astype(np.float32)).unsqueeze(0)
return self.model(img, words, s, pucks)[0].numpy()
def act(self, t: int, image: np.ndarray, text: str, state: np.ndarray) -> np.ndarray:
chunk = self.chunk(image, text, state)
if not self.ensemble:
return chunk[0]
self.history = [(t0, c) for t0, c in self.history if t - t0 < len(c)] + [(t, chunk)]
preds = np.array([c[t - t0] for t0, c in self.history])
weights = np.exp(-self.decay * np.arange(len(preds))[::-1]) # older chunks weigh more in ACT; see 22.3
return (weights[:, None] * preds).sum(0) / weights.sum()
def run_episode(policy: VLAPolicy, text: str, task, layout: tt.Layout, variation: tt.Variation | None = None, camera: str = “top”, record: bool = False) -> dict: scene = tt.Tabletop(variation) scene.reset(layout) policy.reset() log = {“states”: [], “actions”: [], “images”: []} for t in range(tt.EPISODE_STEPS): s = scene.state() image = scene.render(camera, 64) a = policy.act(t, image, text, s) if record: log[“states”].append(s) log[“actions”].append(a) log[“images”].append(image) scene.step(a) if scene.success(task) and np.linalg.norm(scene.state()[4 + 4 * tt.OBJECTS.index(task[0]) + 2:][:2]) < 0.02: break out = {“success”: scene.success(task), “steps”: t + 1, “final”: scene.state(), **({“log”: log} if record else {})} scene.close() return out
Output of `--vision patch --seed 0` (16 cores of a compute node; the run takes 14 minutes to train and 2 to evaluate):
```text
minivla_patch_s0_pos: 2995 successful demonstrations, 169337 frames
step 1000: loss 0.3734 (29 s)
step 2000: loss 0.3488 (58 s)
step 3000: loss 0.2348 (86 s)
step 4000: loss 0.1444 (113 s)
step 5000: loss 0.0908 (141 s)
step 6000: loss 0.0698 (169 s)
step 7000: loss 0.0601 (196 s)
step 8000: loss 0.0520 (223 s)
step 9000: loss 0.0475 (251 s)
step 10000: loss 0.0430 (278 s)
step 11000: loss 0.0389 (305 s)
step 12000: loss 0.0367 (333 s)
step 13000: loss 0.0329 (360 s)
step 14000: loss 0.0315 (387 s)
step 15000: loss 0.0290 (414 s)
step 16000: loss 0.0265 (441 s)
step 17000: loss 0.0250 (469 s)
step 18000: loss 0.0237 (496 s)
step 19000: loss 0.0223 (522 s)
step 20000: loss 0.0211 (550 s)
step 21000: loss 0.0197 (577 s)
step 22000: loss 0.0188 (604 s)
step 23000: loss 0.0181 (632 s)
step 24000: loss 0.0171 (659 s)
step 25000: loss 0.0166 (693 s)
step 26000: loss 0.0164 (725 s)
step 27000: loss 0.0160 (753 s)
step 28000: loss 0.0157 (781 s)
step 29000: loss 0.0157 (808 s)
step 30000: loss 0.0155 (836 s)
{"vision": "patch", "seed": 0, "steps": 30000, "demos": 3000, "state": 2, "train_seconds": 837.2640732079744, "final_loss": 0.015473779505118728, "summary": {"train": 0.95, "paraphrase": 0.8666666666666667, "synonym": 0.016666666666666666, "composition": 0.7}}
The diagnosis:
```python file=examples/l22_4_copycat.py import json import os from pathlib import Path
import numpy as np import torch
from mjcourse import tabletop as tt from mjcourse.vla import data, language, train
RUNS = Path(file).resolve().parents[1] / “runs” / “vla”
def part1() -> None: print(“(1) every trained MiniVLA: final training loss and closed-loop success (train split: 100 held-out layouts; “ “others: 60)”) print(f” {‘vision’:<10}{‘state’:<22}{‘seed’:>5}{‘loss’:>9}{‘train’:>8}{‘paraphrase’:>12}{‘synonym’:>9}{‘composition’:>13}”) for vision in (“patch”, “keypoint”, “oracle”): for suffix, state in ((“”, “position and velocity”), (“pos”, “position”)): for seed in (0, 1, 2): path = RUNS / f”minivla{vision}_s{seed}{suffix}.json” if not path.exists(): continue r = json.loads(path.read_text()) s = r[“summary”] print(f” {vision:<10}{state:<22}{seed:>5}{r[‘final_loss’]:>9.4f}{s[‘train’]:>8.2f}{s[‘paraphrase’]:>12.2f}” f”{s[‘synonym’]:>9.2f}{s[‘composition’]:>13.2f}”)
def predict(model, frames: list[tuple], tokenizer, zero_velocity: bool = False) -> np.ndarray: out = [] with torch.no_grad(): for image, text, state in frames: img = torch.as_tensor(image).permute(2, 0, 1).float().unsqueeze(0) / 255.0 words = torch.as_tensor(tokenizer(text)).unsqueeze(0) s = (state[:4] * train.STATE_SCALE).astype(np.float32) if zero_velocity: s = s * np.array([1, 1, 0, 0], np.float32) out.append(model(img, words, torch.as_tensor(s[:model.cfg.state_dim]).unsqueeze(0))[0, 0].numpy()) return np.array(out)
def part2_3(models: dict, episodes: list[dict], tokenizer) -> None: first = [(e[“images”][0], e[“text”], e[“states”][0]) for e in episodes] first_target = np.array([e[“actions”][0] for e in episodes]) later = [(e, t) for e in episodes for t in range(1, len(e[“actions”]))] later_frames = [(e[“images”][t], e[“text”], e[“states”][t]) for e, t in later] later_target = np.array([e[“actions”][t] for e, t in later]) copy_error = np.linalg.norm(later_target - np.array([e[“actions”][t - 1] for e, t in later]), axis=1).mean() print(f”(2) offline action error on {len(episodes)} held-out demonstrations (mean distance; demonstrator’s actions “ f”have mean norm {np.linalg.norm(later_target, axis=1).mean():.2f})”) print(f” {‘model’:<34}{‘first frame (at rest)’:>22}{‘later frames’:>14}”) for name, model in models.items(): e0 = np.linalg.norm(predict(model, first, tokenizer) - first_target, axis=1).mean() e1 = np.linalg.norm(predict(model, later_frames, tokenizer) - later_target, axis=1).mean() print(f” {name:<34}{e0:>22.3f}{e1:>14.3f}”) print(f” {‘copy the previous action’:<34}{‘(none)’:>22}{copy_error:>14.3f}”) model = models[“patch, position and velocity”] with_v = predict(model, later_frames, tokenizer) without_v = predict(model, later_frames, tokenizer, zero_velocity=True) velocity = np.array([s[2:4] for _, _, s in later_frames]) / tt.MAX_SPEED # the velocity, as an action print(f”(3) the velocity model on the {len(later_frames)} later frames”) print(f” action change when the velocity input is set to zero: mean {np.linalg.norm(with_v - without_v, axis=1).mean():.3f}”) print(f” distance between its action and the velocity it is given (in action units): “ f”{np.linalg.norm(with_v - velocity, axis=1).mean():.3f}; between the demonstrator’s action and that “ f”velocity: {np.linalg.norm(later_target - velocity, axis=1).mean():.3f}”)
def part4(paths: dict, tokenizer, tasks: int = 20, seconds: int = 3) -> None: print(f”(4) closed loop from rest, {tasks} held-out tasks: distance between the policy’s action and the demonstrator’s “ f”action at the states the policy itself visits, mean per second”) print(f” {‘model’:<34}” + ““.join(f”{f’{k}-{k + 1} s’:>10}” for k in range(seconds))) for name, path in paths.items(): model = torch.load(path, weights_only=False) errors = np.zeros((tasks, 10 * seconds)) for i in range(tasks): rng = np.random.default_rng([2027, i]) text, task = language.instructions(“train”)[rng.integers(len(language.instructions(“train”)))] scene = tt.Tabletop() scene.reset(tt.sample_layout(rng, task)) policy = train.VLAPolicy(model, tokenizer) for t in range(10 * seconds): s = scene.state() a = policy.act(t, scene.render(“top”, 64), text, s) errors[i, t] = np.linalg.norm(a - tt.expert(s, task)) scene.step(a) scene.close() per_second = errors.reshape(tasks, seconds, 10).mean((0, 2)) print(f” {name:<34}” + ““.join(f”{e:>10.3f}” for e in per_second))
def figure(paths: dict, tokenizer, index: int = 0) -> None: import matplotlib matplotlib.use(“Agg”) import matplotlib.pyplot as plt
rng = np.random.default_rng([2027, index])
text, task = language.instructions("train")[rng.integers(len(language.instructions("train")))]
layout = tt.sample_layout(rng, task)
runs = {"demonstrator": None, **paths}
fig, axes = plt.subplots(1, 3, figsize=(10, 3.7), constrained_layout=True, sharex=True, sharey=True)
colours = {"red": "#D55E00", "green": "#009E73", "blue": "#0072B2"}
for ax, (name, path) in zip(axes, runs.items()):
model = torch.load(path, weights_only=False) if path else None
policy = train.VLAPolicy(model, tokenizer) if model else None
scene = tt.Tabletop()
scene.reset(layout)
states = [scene.state()]
for t in range(60):
s = scene.state()
scene.step(policy.act(t, scene.render("top", 64), text, s) if policy else tt.expert(s, task))
states.append(scene.state())
scene.close()
states = np.array(states)
for zone, (zx, zy) in tt.ZONES.items():
ax.add_patch(plt.Circle((zy, zx), tt.SUCCESS_RADIUS, color="#bbbbbb", alpha=0.5 if zone == task[1] else 0.2))
for k, obj in enumerate(tt.OBJECTS):
i = 4 + 4 * k
ax.plot(states[:, i + 1], states[:, i], color=colours[obj], lw=2 if obj == task[0] else 1)
ax.add_patch(plt.Circle((states[0, i + 1], states[0, i]), tt.PUCK_RADIUS, fill=False, color=colours[obj], ls=":"))
ax.plot(states[:, 1], states[:, 0], color="black", lw=1)
ax.plot(states[0, 1], states[0, 0], "ko", ms=4)
ax.set_title(name, fontsize=10)
ax.set_aspect("equal")
ax.set_xlim(0.3, -0.3) # as the top camera shows it: +y to the left
ax.set_ylim(-0.3, 0.3)
ax.set_xlabel("y (m)")
axes[0].set_ylabel("x (m)")
fig.suptitle(f'"{text}", 6 s from the same start (black: pusher; dotted: pucks at the start)', fontsize=10)
fig.savefig(RUNS / "copycat.png", dpi=150)
plt.close(fig)
if name == “main”: torch.set_num_threads(min(8, os.cpu_count() or 1)) tokenizer = language.default_tokenizer() part1() paths = {“patch, position and velocity”: RUNS / “minivla_patch_s0.pt”, “patch, position only”: RUNS / “minivla_patch_s0_pos.pt”} models = {k: torch.load(p, weights_only=False).eval() for k, p in paths.items()} episodes = [e for e in data.generate(7, 40) if e[“success”]] part2_3(models, episodes, tokenizer) part4(paths, tokenizer) figure(paths, tokenizer)
```text
(1) every trained MiniVLA: final training loss and closed-loop success (train split: 100 held-out layouts; others: 60)
vision state seed loss train paraphrase synonym composition
patch position and velocity 0 0.0173 0.01 0.00 0.00 0.00
patch position 0 0.0155 0.95 0.87 0.02 0.70
patch position 1 0.0151 0.92 0.77 0.32 0.82
patch position 2 0.0142 0.99 0.83 0.23 0.95
keypoint position and velocity 0 0.0179 0.01 0.02 0.02 0.00
keypoint position 0 0.0315 0.58 0.02 0.02 0.08
keypoint position 1 0.0111 0.96 0.23 0.00 0.72
keypoint position 2 0.0311 0.62 0.03 0.00 0.10
oracle position and velocity 0 0.0160 0.01 0.00 0.00 0.00
oracle position 0 0.0068 0.99 0.23 0.07 0.63
oracle position 1 0.0047 1.00 0.92 0.32 0.85
oracle position 2 0.0067 1.00 0.00 0.07 0.62
(2) offline action error on 40 held-out demonstrations (mean distance; demonstrator's actions have mean norm 0.58)
model first frame (at rest) later frames
patch, position and velocity 0.653 0.058
patch, position only 0.079 0.090
copy the previous action (none) 0.071
(3) the velocity model on the 2313 later frames
action change when the velocity input is set to zero: mean 0.752
distance between its action and the velocity it is given (in action units): 0.072; between the demonstrator's action and that velocity: 0.092
(4) closed loop from rest, 20 held-out tasks: distance between the policy's action and the demonstrator's action at the states the policy itself visits, mean per second
model 0-1 s 1-2 s 2-3 s
patch, position and velocity 0.739 1.109 1.300
patch, position only 0.129 0.163 0.184
The demonstrator decides what can be imitated. The tabletop’s first demonstrator switched between an approach mode and a push mode at a distance threshold. Its actions jumped at the switch, and policies cloned from its state did poorly; collecting corrections with DAgger made them worse, because the corrections near the threshold contradicted each other. The demonstrator used here is one continuous velocity field: approach the point behind the puck, slide around it, push. Check a demonstrator’s imitability before training a VLA on it: clone it from the true state first, where perception cannot be the problem.
Low loss, no success. Lower loss means a closer fit to the demonstration frames, which the policy will not visit once its own errors move it elsewhere. Evaluate in closed loop, and when closed loop fails, compare the first-frame error and the on-policy disagreement (outputs 2 and 4) with the held-out loss.
One seed lied twice. Keypoint seed 0 alone says the encoder reaches 0.58; seed 1 alone says 0.96. Three seeds are the minimum before saying anything about an encoder, and the spread across seeds belongs in the result.
Train the patch model with 1000 demonstrations instead of 3000 (--demos 1000, seeds 0 to 2) and report success on training instructions with the seed spread. Predict whether the loss or the success drops more, and why.
Keep the velocity in the input but break the shortcut: during training, replace the velocity with the velocity recorded $k$ frames earlier, or add noise of the size of the demonstrator’s action changes. Find the smallest change that restores closed-loop success above 0.9, and check outputs (2) to (4) for the resulting model.
[!research] Shortcuts in large VLAs Large VLAs differ from the MiniVLA in what they take in: OpenVLA (arXiv:2406.09246) conditions on a single image and the instruction, with no proprioceptive history, while π0 (arXiv:2410.24164) takes the robot’s proprioceptive state. Whether a given input lets a large model take a shortcut like this one is an empirical question about that model and its data; the papers in Lesson 22.1’s table do not report the first-frame and on-policy measurements used here, so their susceptibility is [unverified]. The tests of outputs (2) to (4) apply to any policy whose actions can be queried at arbitrary states.
The planar pusher isolates causal confusion in its purest form. Real robot deployments, however, operate multi-joint manipulators. The course companion code provides a full 7-DOF Franka Emika Panda simulation in MuJoCo (franka_vla.xml and mjcourse.vla.franka.FrankaVLA).
The 7-DOF Franka model incorporates:
gripper/wrist).top) paired with egocentric visual feedback from the wrist camera.The causal confusion failure observed on the planar pusher translates directly to the Franka manipulator. If joint velocities $\dot{q} \in \mathbb{R}^7$ or TCP velocities are provided in the proprioceptive state token, the policy learns to integrate velocity commands rather than grounding the target objects in the image tokens. Restricting the Franka proprioceptive vector to position-only quantities (TCP coordinates $p \in \mathbb{R}^3$, orientation quaternion $q \in \mathbb{H}$, gripper width $g \in \mathbb{R}$, and joint angles $q_{\text{arm}} \in \mathbb{R}^7$, 15 dimensions total) forces the policy to ground its actions in the visual and language tokens.
```python file=examples/l22_4_franka_vla.py from future import annotations
import os import numpy as np
from mjcourse.vla.franka import FrankaVLA, franka_expert_action
FAST = os.environ.get(“MJC_FAST”) == “1” STEPS = 20 if FAST else 120
def main() -> None: print(“Initializing 7-DOF Franka Emika Panda VLA simulator in MuJoCo…”) env = FrankaVLA()
# (1) Topology
nq, nv, nu = env.model.nq, env.model.nv, env.model.nu
print(f"Topology: nq = {nq}, nv = {nv}, nu = {nu} actuators (7 arm torques + 1 gripper)")
# (2) Observations
pos, quat = env.tcp_pose()
state_no_vel = env.state(include_velocity=False)
state_with_vel = env.state(include_velocity=True)
print(f"TCP Pose: position = {np.round(pos, 3)}, quaternion = {np.round(quat, 3)}")
print(f"Proprioception: {len(state_no_vel)} dims (position-only) vs {len(state_with_vel)} dims (with velocity)")
# (3) Dual-camera rendering
img_top = env.render("top", width=64, height=64)
img_wrist = env.render("gripper/wrist", width=64, height=64)
print(f"Rendered views: top = {img_top.shape}, wrist = {img_wrist.shape}")
# (4) Rollout
task = ("red_cube", "tray")
env.reset(seed=42, task=task)
print(f"Task: move {task[0]} to {task[1]}")
success = False
for t in range(STEPS):
action = franka_expert_action(env)
step_result = env.step(action)
if step_result["success"]:
success = True
print(f"Goal achieved at step {t} (distance to tray = {step_result['dist_to_target'] * 1000:.1f} mm)")
break
if not success:
print(f"Rollout finished: final distance to target = {step_result['dist_to_target'] * 1000:.1f} mm")
env.close()
print("7-DOF Franka VLA simulation interface verified.")
if name == “main”: main()
```quiz
{"id": "22.4-check", "title": "Knowledge check", "questions": [
{"kind": "mcq", "q": "The oracle MiniVLA, which reads true puck positions instead of an image, also failed with the velocity input. What did that rule out?",
"options": ["The demonstrator", "Perception as the cause", "The transformer", "The language encoder"],
"answer": 1,
"explain": "<p>A failure that survives perfect perception cannot be a perception failure.</p>"},
{"kind": "mcq", "q": "Why can a network predict the demonstrator's action from the pusher's velocity alone?",
"options": ["Velocity encodes the puck's position", "Velocity actuators track the previous command, and the commands change smoothly, so the velocity is almost the next action", "The demonstrator ignores the pucks", "Velocity is normalized"],
"answer": 1,
"explain": "<p>Copying the previous action already has an error of 0.071.</p>"},
{"kind": "mcq", "q": "The velocity model's offline error on later frames is 0.058. Its disagreement with the demonstrator in closed loop during the first second is:",
"options": ["0.058", "0.13", "0.74", "0.0"],
"answer": 2,
"explain": "<p>In closed loop its velocity input is its own previous action, so errors persist and grow (1.11, then 1.30).</p>"},
{"kind": "mcq", "q": "Which measurement would have caught the shortcut before any closed-loop run?",
"options": ["The final training loss", "The error at each episode's first frame, where the velocity is zero", "The parameter count", "The training time"],
"answer": 1,
"explain": "<p>0.653 at rest against 0.058 elsewhere: the model has not learned the task without its velocity.</p>"}
]}
Lesson 22.5 evaluates the position-only patch model along one generalization axis at a time: language, objects, space, camera and embodiment, with three seeds and intervals.